Update app.py

app.py

@@ -67,10 +67,18 @@ def img_to_img_fn(model_str, image, prompt, negative_prompt, max_retries=10):
     while retries < max_retries:
         try:
             noise = str(randint(0, 9999999))
-            if hasattr(models_load[model_str], 'negative_prompt'):
-                result = models_load[model_str](image=image, prompt=f'{prompt} {noise}', negative_prompt=negative_prompt)
+            if hasattr(models_load[model_str], 'img2img'):
+                # Assuming the model has an 'img2img' method
+                result = models_load[model_str].img2img(image=image, prompt=f'{prompt} {noise}', negative_prompt=negative_prompt)
+            elif hasattr(models_load[model_str], 'image_to_image'):
+                # Assuming the model has an 'image_to_image' method
+                result = models_load[model_str].image_to_image(image=image, prompt=f'{prompt} {noise}', negative_prompt=negative_prompt)
+            elif hasattr(models_load[model_str], 'image2image'):
+                # Assuming the model has an 'image2image' method
+                result = models_load[model_str].image2image(image=image, prompt=f'{prompt} {noise}', negative_prompt=negative_prompt)
             else:
-                result = models_load[model_str](image=image, prompt=f'{prompt} {noise}')
+                # Fallback to a generic method
+                result = models_load[model_str](image=image, prompt=f'{prompt} {noise}', negative_prompt=negative_prompt)
             return result
         except Exception as e:
             # Check for specific error messages or status codes