Ming Li commited on
Commit
3029694
1 Parent(s): 01df415

add Decorator for each condition

Browse files
Files changed (1) hide show
  1. model.py +5 -0
model.py CHANGED
@@ -126,6 +126,7 @@ class Model:
126
  ).images
127
 
128
  @torch.inference_mode()
 
129
  def process_canny(
130
  self,
131
  image: np.ndarray,
@@ -304,6 +305,7 @@ class Model:
304
  return [control_image] + results
305
 
306
  @torch.inference_mode()
 
307
  def process_softedge(
308
  self,
309
  image: np.ndarray,
@@ -416,6 +418,7 @@ class Model:
416
  return [control_image] + results
417
 
418
  @torch.inference_mode()
 
419
  def process_segmentation(
420
  self,
421
  image: np.ndarray,
@@ -469,6 +472,7 @@ class Model:
469
  return [control_image] * num_images + results + conditions_of_generated_imgs
470
 
471
  @torch.inference_mode()
 
472
  def process_depth(
473
  self,
474
  image: np.ndarray,
@@ -566,6 +570,7 @@ class Model:
566
  return [control_image] + results
567
 
568
  @torch.inference_mode()
 
569
  def process_lineart(
570
  self,
571
  image: np.ndarray,
 
126
  ).images
127
 
128
  @torch.inference_mode()
129
+ @spaces.GPU()
130
  def process_canny(
131
  self,
132
  image: np.ndarray,
 
305
  return [control_image] + results
306
 
307
  @torch.inference_mode()
308
+ @spaces.GPU()
309
  def process_softedge(
310
  self,
311
  image: np.ndarray,
 
418
  return [control_image] + results
419
 
420
  @torch.inference_mode()
421
+ @spaces.GPU()
422
  def process_segmentation(
423
  self,
424
  image: np.ndarray,
 
472
  return [control_image] * num_images + results + conditions_of_generated_imgs
473
 
474
  @torch.inference_mode()
475
+ @spaces.GPU()
476
  def process_depth(
477
  self,
478
  image: np.ndarray,
 
570
  return [control_image] + results
571
 
572
  @torch.inference_mode()
573
+ @spaces.GPU()
574
  def process_lineart(
575
  self,
576
  image: np.ndarray,