Update app.py

app.py

@@ -1,10 +1,9 @@
 import os
-import gradio as gr
 from random import randint
-from diffusers import DiffusionPipeline
 from flask import Flask, request, send_file
 from io import BytesIO
 from PIL import Image, ImageChops
+from diffusers import DiffusionPipeline
 
 app = Flask(__name__)
 
@@ -15,10 +14,7 @@ def load_model(model_name):
     global models_load
     if model_name not in models_load:
         try:
-            # Assume models are hosted on the HuggingFace model hub
-            model = DiffusionPipeline.from_pretrained(model_name)
-            model = model.to("cuda")
-            models_load[model_name] = model
+            models_load[model_name] = DiffusionPipeline.from_pretrained(f'models/{model_name}')
         except Exception as error:
             print(f"Error loading model {model_name}: {error}")
             models_load[model_name] = None
@@ -27,46 +23,42 @@ def gen_fn(model_str, prompt, negative_prompt=None, noise=None, cfg_scale=None,
     if model_str not in models_load:
         load_model(model_str)  # モデルがロードされていない場合はロードする
 
-    if model_str in models_load:
-        model = models_load[model_str]
-        if model is None:
-            return None, f"Model {model_str} could not be loaded."
+    model = models_load.get(model_str)
+    if not model:
+        print(f"Model {model_str} not found or failed to load.")
+        return None, f"Model {model_str} not found or failed to load."
 
-        if noise == "random":
-            noise = str(randint(0, 99999999999))
-        full_prompt = f'{prompt} {noise}' if noise else prompt
+    if noise == "random":
+        noise = str(randint(0, 99999999999))
+    full_prompt = f'{prompt} {noise}' if noise else prompt
 
-        try:
-            # Generate the image with the model
-            if cfg_scale is not None and num_inference_steps is not None:
-                result = model(full_prompt, negative_prompt=negative_prompt, cfg_scale=cfg_scale, num_inference_steps=num_inference_steps).images[0]
-            elif cfg_scale is not None:
-                result = model(full_prompt, negative_prompt=negative_prompt, cfg_scale=cfg_scale).images[0]
-            elif num_inference_steps is not None:
-                result = model(full_prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps).images[0]
-            else:
-                result = model(full_prompt, negative_prompt=negative_prompt).images[0]
+    try:
+        # Generate the image
+        kwargs = {"prompt": full_prompt}
+        if negative_prompt:
+            kwargs["negative_prompt"] = negative_prompt
+        if cfg_scale is not None:
+            kwargs["guidance_scale"] = cfg_scale
+        if num_inference_steps is not None:
+            kwargs["num_inference_steps"] = num_inference_steps
+
+        result = model(**kwargs).images[0]
 
-            # Ensure the result is an image
-            if isinstance(result, Image.Image):
-                image = result
-            else:
-                print("Result is not an image:", type(result))
-                return None, 'Result is not an image'
-            
+        # Check if result is an image
+        if isinstance(result, Image.Image):
             # Check if the image is completely black
-            black = Image.new('RGB', image.size, (0, 0, 0))
-            if ImageChops.difference(image, black).getbbox() is None:
+            black = Image.new('RGB', result.size, (0, 0, 0))
+            if ImageChops.difference(result, black).getbbox() is None:
                 return None, 'The image is completely black. There may be a parameter that cannot be specified, or an error may have occurred internally.'
+            return result, None
 
-            return image, None
+        else:
+            print("Result is not an image:", type(result))
+            return None, 'Result is not an image'
 
-        except Exception as e:
-            print("Error generating image:", e)
-            return None, f"Error generating image: {e}"
-    else:
-        print(f"Model {model_str} not found")
-        return None, f"Model {model_str} not found"
+    except Exception as e:
+        print("Error generating image:", e)
+        return None, f"Error generating image: {e}"
 
 @app.route('/', methods=['GET'])
 def home():