tianleliphoebe commited on
Commit
b4c89cc
1 Parent(s): 8752733

Update model/model_manager.py

Browse files
Files changed (1) hide show
  1. model/model_manager.py +6 -1
model/model_manager.py CHANGED
@@ -54,7 +54,12 @@ class ModelManager:
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)
 
 
 
 
 
58
  return result
59
 
60
  def generate_image_ie_parallel(self, textbox_source, textbox_target, textbox_instruct, source_image, model_A, model_B):
 
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
+ if 'PNP' in model_name:
58
+ result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct, num_inversion_steps=300)
59
+ elif 'Prompt2prompt' in model_name:
60
+ result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct, num_inner_steps=5)
61
+ else:
62
+ result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct)
63
  return result
64
 
65
  def generate_image_ie_parallel(self, textbox_source, textbox_target, textbox_instruct, source_image, model_A, model_B):