Update app.py

app.py

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