Update app.py

app.py

@@ -168,36 +168,33 @@ def process_image(img, width, height):
     return img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)
 
 
-def txt_to_img(model_path, prompt, n_images, neg_prompt, guidance, steps, width, height, generator, seed):
-    print(f"{datetime.datetime.now()} txt_to_img, model: {current_model.name}")
-
-    load_model(model_path, "txt2img")
+def inference_image(pipe, model_path, prompt, mode, *args, **kwargs):
+    print(f"{datetime.datetime.now()} {mode}, model: {model_path}")
+    load_model(pipe, model_path, mode)
     prompt = current_model.prefix + prompt
-
     result = pipe(
         prompt,
-        negative_prompt=neg_prompt,
+        *args,
+        **kwargs,
+        callback=pipe_callback
+    )
+    return replace_nsfw_images(result)
+
+def txt_to_img(pipe, model_path, prompt, n_images, neg_prompt, guidance, steps, width, height, generator, seed):
+    return inference_image(pipe, model_path, prompt, "txt2img", 
+        negative_prompt=neg_prompt, 
         num_images_per_prompt=n_images,
         num_inference_steps=int(steps),
         guidance_scale=guidance,
         width=width,
         height=height,
         generator=generator,
-        callback=pipe_callback
+        seed=seed
     )
 
-    return replace_nsfw_images(result)
-
-
-def img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance, steps, width, height, generator, seed):
-    print(f"{datetime.datetime.now()} img_to_img, model: {model_path}")
-
-    load_model(model_path, "img2img")
-    prompt = current_model.prefix + prompt
+def img_to_img(pipe, model_path, prompt, n_images, neg_prompt, img, strength, guidance, steps, width, height, generator, seed):
     img = process_image(img, width, height)
-
-    result = pipe(
-        prompt,
+    return inference_image(pipe, model_path, prompt, "img2img", 
         negative_prompt=neg_prompt,
         num_images_per_prompt=n_images,
         image=img,
@@ -205,11 +202,9 @@ def img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance
         strength=strength,
         guidance_scale=guidance,
         generator=generator,
-        callback=pipe_callback
+        seed=seed
     )
 
-    return replace_nsfw_images(result)
-
 def replace_nsfw_images(results):
 
     if is_colab: