nick_93 commited on
Commit
60efe0d
1 Parent(s): b47e7c8
Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -13,6 +13,7 @@ from utils import map_colors_rgb
13
  from diffusers import StableDiffusionXLPipeline
14
  import gradio as gr
15
  import gc
 
16
 
17
  device = "cuda"
18
  dtype = torch.float16
@@ -209,7 +210,8 @@ class ControlNetDepthDesignModelMulti:
209
  self.seg_image_processor, self.image_segmentor = get_segmentation_pipeline()
210
  self.depth_feature_extractor, self.depth_estimator = get_depth_pipeline()
211
  self.depth_estimator = self.depth_estimator.to(device)
212
-
 
213
  def generate_design(self, empty_room_image: Image, prompt: str, guidance_scale: int = 10, num_steps: int = 50, strength: float =0.9, img_size: int = 640) -> Image:
214
  """
215
  Given an image of an empty room and a prompt
 
13
  from diffusers import StableDiffusionXLPipeline
14
  import gradio as gr
15
  import gc
16
+ import spaces
17
 
18
  device = "cuda"
19
  dtype = torch.float16
 
210
  self.seg_image_processor, self.image_segmentor = get_segmentation_pipeline()
211
  self.depth_feature_extractor, self.depth_estimator = get_depth_pipeline()
212
  self.depth_estimator = self.depth_estimator.to(device)
213
+
214
+ @spaces.GPU
215
  def generate_design(self, empty_room_image: Image, prompt: str, guidance_scale: int = 10, num_steps: int = 50, strength: float =0.9, img_size: int = 640) -> Image:
216
  """
217
  Given an image of an empty room and a prompt