add config of init_image.

mazzzystar · May 5, 2022 · 148262c · 148262c
1 parent 6fb676b
commit 148262c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/wrapper/run.py b/wrapper/run.py
@@ -42,7 +42,7 @@ def load_diffusion_model(batch_name):
     return model, diffusion, clip_models, secondary_model, lpips_model
 
 
-def diffuse(text_prompt, batch_name, image_prompts={}, steps=200, batch_size=1, display_rate=40):
+def diffuse(text_prompt, batch_name, image_prompts={}, init_image=None, steps=200, batch_size=1, display_rate=40):
     model, diffusion, clip_models, secondary_model, lpips_model = load_diffusion_model(batch_name)
     global args
     batchNum = args["batchNum"]
@@ -57,6 +57,7 @@ def diffuse(text_prompt, batch_name, image_prompts={}, steps=200, batch_size=1,
     args['steps'] = steps
     args['batch_size'] = batch_size
     args['display_rate'] = display_rate
+    args['init_image'] = init_image
 
     args = SimpleNamespace(**args)
 
@@ -76,11 +77,11 @@ def diffuse(text_prompt, batch_name, image_prompts={}, steps=200, batch_size=1,
 if __name__ == '__main__':
     print(str(sys.argv[1]))
     input_text = str(sys.argv[1])
-    input_img_path = "" # If do not need input img，then set None or ""
+    input_img_path = None # If do not need input img，then set None or ""
     result = translate(input_text)
     del translator
     text_prompts, image_prompts = simple_prompts(result, input_img_path)
     print(text_prompts)
     outdirName = "my-test"
-    diffuse(text_prompts, outdirName, steps=240, image_prompts=image_prompts, display_rate=40)
+    diffuse(text_prompts, outdirName, steps=240, image_prompts=image_prompts, init_image=input_img_path, display_rate=40)
 
diff --git a/wrapper/run_batch.py b/wrapper/run_batch.py
@@ -64,7 +64,7 @@ def __init__(self, batch_name):
         self.batch_name = batch_name
         self.model, self.diffusion, self.clip_models, self.secondary_model, self.lpips_model = load_diffusion_model(batch_name)
 
-    def draw(self, text_prompts, orig_text, image_prompts={}, steps=200, batch_size=1, display_rate=40):
+    def draw(self, text_prompts, orig_text, image_prompts={}, init_image=None, steps=200, batch_size=1, display_rate=40):
         global args
         batchNum = args["batchNum"]
         start_frame = args["start_frame"]
@@ -78,6 +78,7 @@ def draw(self, text_prompts, orig_text, image_prompts={}, steps=200, batch_size=
         args['steps'] = steps
         args['batch_size'] = batch_size
         args['display_rate'] = display_rate
+        args['init_image'] = init_image
 
         args = SimpleNamespace(**args)
 
@@ -114,12 +115,12 @@ def draw(self, text_prompts, orig_text, image_prompts={}, steps=200, batch_size=
                 from utils import *
                 input_text = line.strip() + "，" + str(style)
                 print(input_text)
-                input_img_path = "" # If image input not needed, set None or ""
+                input_img_path = None # If image input not needed, set None or ""
                 result = translate(input_text)
 
                 text_prompts, image_prompts = simple_prompts(result, input_img_path)
                 print(text_prompts)
-                disco.draw(text_prompts, input_text, steps=240, image_prompts=image_prompts, display_rate=40)
+                disco.draw(text_prompts, input_text, steps=240, image_prompts=image_prompts, init_image=input_img_path, display_rate=40)
             except Exception as e:
                 print(e)
     del translator