Update app.py

app.py

@@ -27,7 +27,7 @@ def load_model(model_name):
             print(f"Error loading model {model_name}: {error}")
             models_load[model_name] = gr.Interface(lambda txt: None, ['text'], ['image'])
 
-def gen_fn(model_str, prompt, negative_prompt=None, noise=None, cfg_scale=None):
+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)  # モデルがロードされていない場合はロードする
 
@@ -38,11 +38,16 @@ def gen_fn(model_str, prompt, negative_prompt=None, noise=None, cfg_scale=None):
         try:
             if negative_prompt:
                 full_prompt += f' -{negative_prompt}'
-            # Adjust function call based on whether cfg_scale is provided
+            
+            # Construct the function call parameters dynamically
+            call_params = {'text': full_prompt}
             if cfg_scale is not None:
-                result = models_load[model_str](full_prompt, cfg_scale=cfg_scale)
-            else:
-                result = models_load[model_str](full_prompt)
+                call_params['cfg_scale'] = cfg_scale
+            if num_inference_steps is not None:
+                call_params['num_inference_steps'] = num_inference_steps
+            
+            result = models_load[model_str](**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):
@@ -69,12 +74,19 @@ def home():
     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)
     
     try:
         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:
+        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
@@ -83,7 +95,7 @@ def home():
         return 'Please provide a "prompt" query parameter in the URL.', 400
     
     # Generate the image
-    image = gen_fn(model, prompt, negative_prompt, noise, cfg_scale)
+    image = gen_fn(model, prompt, negative_prompt, noise, cfg_scale, num_inference_steps)
     if isinstance(image, Image.Image):  # Ensure the result is a PIL image
         # Save image to BytesIO object
         img_io = BytesIO()