Update app.py

app.py

@@ -26,102 +26,56 @@ def load_model(model_name):
             models_load[model_name] = m
         except Exception as error:
             print(f"Error loading model {model_name}: {error}")
-            models_load[model_name] = gr.Interface(lambda txt: None, ['text'], ['image'])
+            models_load[model_name] = None  # Set to None if loading fails
 
 def gen_fn(model_str, prompt, negative_prompt=None, noise=None, cfg_scale=None, num_inference_steps=None):
     if model_str not in models_load:
         load_model(model_str)  # モデルがロードされていない場合はロードする
 
-    if model_str in models_load:
-        if noise == "random":
-            noise = str(randint(0, 99999999999))
-        full_prompt = f'{prompt} {noise}' if noise else prompt
-        try:
-            # Construct the function call parameters dynamically
-            call_params = {'prompt': full_prompt}
-            if negative_prompt:
-                call_params['negative_prompt'] = negative_prompt
-            if cfg_scale is not None:
-                # Check if cfg_scale can be set
-                if 'cfg_scale' not in models_load[model_str].parameters:
-                    return None, 'cfg_scale parameter is not supported by the model'
-                call_params['cfg_scale'] = cfg_scale
-            if num_inference_steps is not None:
-                # Check if num_inference_steps can be set
-                if 'num_inference_steps' not in models_load[model_str].parameters:
-                    return None, 'num_inference_steps parameter is not supported by the model'
-                call_params['num_inference_steps'] = num_inference_steps
-            
-            result = models_load[model_str].predict(**call_params)
-            
-            # Check if result is an image or a file path
-            if isinstance(result, str):  # Assuming result might be a file path
-                if os.path.exists(result):
-                    image = Image.open(result)
-                else:
-                    print(f"File path not found: {result}")
-                    return None, 'File path not found'
-            elif isinstance(result, Image.Image):
-                image = result
-            else:
-                print("Result is not an image:", type(result))
-                return None, 'Result is not an 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:
-                return None, 'The image is completely black. There may be a parameter that cannot be specified, or an error may have occurred internally.'
-
-            return image, None
-
-        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"
+    model = models_load.get(model_str)
+    if model is None:
+        return None, f"Model {model_str} not found or failed to load"
 
-@app.route('/', methods=['GET'])
-def home():
-    prompt = request.args.get('prompt', '')
-    model = request.args.get('model', '')
-    negative_prompt = request.args.get('Nprompt', None)
-    noise = request.args.get('noise', None)
-    cfg_scale = request.args.get('cfg_scale', None)
-    num_inference_steps = request.args.get('steps', None)
-    
+    if noise == "random":
+        noise = str(randint(0, 99999999999))
+    full_prompt = f'{prompt} {noise}' if noise else prompt
     try:
+        # Construct the function call parameters dynamically
+        call_params = {'prompt': full_prompt}
+        if negative_prompt:
+            call_params['negative_prompt'] = negative_prompt
         if cfg_scale is not None:
-            cfg_scale = float(cfg_scale)
-    except ValueError:
-        return 'Invalid "cfg_scale" parameter. It should be a number.', 400
-
-    try:
+            # Check if cfg_scale can be set
+            if 'cfg_scale' not in model.parameters:
+                return None, 'cfg_scale parameter is not supported by the model'
+            call_params['cfg_scale'] = cfg_scale
         if num_inference_steps is not None:
-            num_inference_steps = int(num_inference_steps)
-    except ValueError:
-        return 'Invalid "steps" parameter. It should be an integer.', 400
-    
-    if not model:
-        return 'Please provide a "model" query parameter in the URL.', 400
+            # Check if num_inference_steps can be set
+            if 'num_inference_steps' not in model.parameters:
+                return None, 'num_inference_steps parameter is not supported by the model'
+            call_params['num_inference_steps'] = num_inference_steps
+        
+        # Ensure that model's call method is used correctly
+        if hasattr(model, 'predict'):
+            result = model.predict(**call_params)
+        elif hasattr(model, 'process'):
+            result = model.process(**call_params)
+        else:
+            return None, f"Model {model_str} does not have a callable method"
 
-    if not prompt:
-        return 'Please provide a "prompt" query parameter in the URL.', 400
-    
-    # Generate the image
-    image, error_message = gen_fn(model, prompt, negative_prompt, noise, cfg_scale, num_inference_steps)
-    if error_message:
-        return error_message, 400
-    
-    if isinstance(image, Image.Image):  # Ensure the result is a PIL image
-        # Save image to BytesIO object
-        img_io = BytesIO()
-        image.save(img_io, format='PNG')
-        img_io.seek(0)
-        return send_file(img_io, mimetype='image/png', as_attachment=False)
-    
-    return 'Failed to generate image.', 500
-
-if __name__ == '__main__':
-    # Launch Flask app
-    app.run(host='0.0.0.0', port=7860)  # Run Flask app
+        # Check if result is an image or a file path
+        if isinstance(result, str):  # Assuming result might be a file path
+            if os.path.exists(result):
+                image = Image.open(result)
+            else:
+                print(f"File path not found: {result}")
+                return None, 'File path not found'
+        elif isinstance(result, Image.Image):
+            image = result
+        else:
+            print("Result is not an image:", type(result))
+            return None, 'Result is not an image'
+        
+        # Check if the image is completely black
+        black = Image.new('RGB', image.size, (0, 0, 0))
+        if ImageChops.difference(image, black).get