Update app.py

app.py

@@ -127,82 +127,113 @@ def inference(model_name, prompt, guidance, steps, n_images=1, width=512, height
 
   try:
     if img is not None:
-      return img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance, steps, width, height, generator), f"Done. Seed: {seed}"
+      return img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance, steps, width, height, generator, seed), f"Done. Seed: {seed}"
     else:
-      return txt_to_img(model_path, prompt, n_images, neg_prompt, guidance, steps, width, height, generator), f"Done. Seed: {seed}"
+      return txt_to_img(model_path, prompt, n_images, neg_prompt, guidance, steps, width, height, generator, seed), f"Done. Seed: {seed}"
   except Exception as e:
     return None, error_str(e)
 
-def load_model(model_path, mode):
+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}")
+
     global last_mode
     global pipe
     global current_model_path
-
-    if model_path != current_model_path or last_mode != mode:
+    if model_path != current_model_path or last_mode != "txt2img":
         current_model_path = model_path
-        update_state(f"Loading {current_model.name} {mode} model...")
 
-        model_class = StableDiffusionPipeline if mode == "txt2img" else StableDiffusionImg2ImgPipeline
+        update_state(f"Loading {current_model.name} text-to-image model...")
+
         if is_colab or current_model == custom_model:
-            pipe = model_class.from_pretrained(
-                current_model_path,
-                torch_dtype=torch.float16,
-                scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler"),
-                safety_checker=lambda images, clip_input: (images, False)
-            )
+          pipe = StableDiffusionPipeline.from_pretrained(
+              current_model_path,
+              torch_dtype=torch.float16,
+              scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler"),
+              safety_checker=lambda images, clip_input: (images, False)
+              )
         else:
-            pipe = model_class.from_pretrained(
-                current_model_path,
-                torch_dtype=torch.float16,
-                scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler")
-            )
+          pipe = StableDiffusionPipeline.from_pretrained(
+              current_model_path,
+              torch_dtype=torch.float16,
+              scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler")
+              )
+          # pipe = pipe.to("cpu")
+          # pipe = current_model.pipe_t2i
 
         if torch.cuda.is_available():
-            pipe = pipe.to("cuda")
-            pipe.enable_xformers_memory_efficient_attention()
+          pipe = pipe.to("cuda")
+          pipe.enable_xformers_memory_efficient_attention()
+        last_mode = "txt2img"
 
-        last_mode = mode
+    prompt = current_model.prefix + prompt  
+    result = pipe(
+      prompt,
+      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)
+
+    # update_state(f"Done. Seed: {seed}")
+    
+    return replace_nsfw_images(result)
 
-def process_image(img, width, height):
-    ratio = min(height / img.height, width / img.width)
-    return img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)
+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}")
 
-def inference_image(model_path, prompt, mode, *args, **kwargs):
-    print(f"{datetime.datetime.now()} {mode}, model: {model_path}")
+    global last_mode
     global pipe
-    load_model(model_path, mode)
+    global current_model_path
+    if model_path != current_model_path or last_mode != "img2img":
+        current_model_path = model_path
+
+        update_state(f"Loading {current_model.name} image-to-image model...")
+
+        if is_colab or current_model == custom_model:
+          pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+              current_model_path,
+              torch_dtype=torch.float16,
+              scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler"),
+              safety_checker=lambda images, clip_input: (images, False)
+              )
+        else:
+          pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+              current_model_path,
+              torch_dtype=torch.float16,
+              scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler")
+              )
+          # pipe = pipe.to("cpu")
+          # pipe = current_model.pipe_i2i
+        
+        if torch.cuda.is_available():
+          pipe = pipe.to("cuda")
+          pipe.enable_xformers_memory_efficient_attention()
+        last_mode = "img2img"
+
     prompt = current_model.prefix + prompt
+    ratio = min(height / img.height, width / img.width)
+    img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)
     result = pipe(
         prompt,
-        *args,
-        **kwargs,
-        callback=pipe_callback
-    )
-    return replace_nsfw_images(result)
-
-def txt_to_img(model_path, prompt, n_images, neg_prompt, guidance, steps, width, height, generator):
-    return inference_image(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
-    )
-
-def img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance, steps, width, height, generator):
-    img = process_image(img, width, height)
-    return inference_image(model_path, prompt, "img2img", 
-        negative_prompt=neg_prompt,
+        negative_prompt = neg_prompt,
         num_images_per_prompt=n_images,
-        image=img,
-        num_inference_steps=int(steps),
-        strength=strength,
-        guidance_scale=guidance,
-        generator=generator
-    )
+        image = img,
+        num_inference_steps = int(steps),
+        strength = strength,
+        guidance_scale = guidance,
+        # width = width,
+        # height = height,
+        generator = generator,
+        callback=pipe_callback)
+
+    # update_state(f"Done. Seed: {seed}")
+        
+    return replace_nsfw_images(result)
 
 def replace_nsfw_images(results):