mikonvergence commited on
Commit
feba4eb
·
1 Parent(s): 625b5c0

Zero GPU test

Browse files
Files changed (1) hide show
  1. src/utils.py +3 -1
src/utils.py CHANGED
@@ -4,9 +4,11 @@ import tempfile
4
  import torch
5
  from scipy.spatial import Delaunay
6
  from .build_pipe import *
 
7
 
8
  pipe = build_pipe()
9
 
 
10
  def generate_terrain(prompt, num_inference_steps, guidance_scale, seed, crop_size, prefix):
11
  """Generates terrain data (RGB and elevation) from a text prompt."""
12
  if prefix and not prefix.endswith(' '):
@@ -56,4 +58,4 @@ def generate_and_display(prompt, num_inference_steps, guidance_scale, seed, crop
56
  mesh.export(temp_file.name)
57
  file_path = temp_file.name
58
 
59
- return file_path
 
4
  import torch
5
  from scipy.spatial import Delaunay
6
  from .build_pipe import *
7
+ import spaces
8
 
9
  pipe = build_pipe()
10
 
11
+ @spaces.GPU
12
  def generate_terrain(prompt, num_inference_steps, guidance_scale, seed, crop_size, prefix):
13
  """Generates terrain data (RGB and elevation) from a text prompt."""
14
  if prefix and not prefix.endswith(' '):
 
58
  mesh.export(temp_file.name)
59
  file_path = temp_file.name
60
 
61
+ return file_path