DongfuJiang commited on
Commit
8752733
1 Parent(s): bd0bce1
Files changed (1) hide show
  1. model/model_manager.py +3 -0
model/model_manager.py CHANGED
@@ -3,6 +3,7 @@ import random
3
  import gradio as gr
4
  import requests
5
  import io, base64, json
 
6
  from PIL import Image
7
  from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, load_pipeline
8
 
@@ -20,6 +21,7 @@ class ModelManager:
20
  pipe = self.loaded_models[model_name]
21
  return pipe
22
 
 
23
  def generate_image_ig(self, prompt, model_name):
24
  pipe = self.load_model_pipe(model_name)
25
  result = pipe(prompt=prompt)
@@ -49,6 +51,7 @@ class ModelManager:
49
  results.append(result)
50
  return results[0], results[1]
51
 
 
52
  def generate_image_ie(self, textbox_source, textbox_target, textbox_instruct, source_image, model_name):
53
  pipe = self.load_model_pipe(model_name)
54
  result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct)
 
3
  import gradio as gr
4
  import requests
5
  import io, base64, json
6
+ import spaces
7
  from PIL import Image
8
  from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, load_pipeline
9
 
 
21
  pipe = self.loaded_models[model_name]
22
  return pipe
23
 
24
+ @spaces.GPU(duration=60)
25
  def generate_image_ig(self, prompt, model_name):
26
  pipe = self.load_model_pipe(model_name)
27
  result = pipe(prompt=prompt)
 
51
  results.append(result)
52
  return results[0], results[1]
53
 
54
+ @spaces.GPU(duration=150)
55
  def generate_image_ie(self, textbox_source, textbox_target, textbox_instruct, source_image, model_name):
56
  pipe = self.load_model_pipe(model_name)
57
  result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct)