Update app.py

app.py

@@ -2,8 +2,6 @@ import gradio as gr
 from random import randint
 from all_models import models
 from datetime import datetime
-from diffusers import StableDiffusionImg2ImgPipeline
-import torch
 
 def get_current_time():
     now = datetime.now()
@@ -18,10 +16,7 @@ def load_fn(models):
     for model in models:
         if model not in models_load.keys():
             try:
-                if 'img2img' in model.lower():  # Assuming models with 'img2img' in their name are image-to-image models
-                    m = StableDiffusionImg2ImgPipeline.from_pretrained(f'models/{model}')
-                else:
-                    m = gr.load(f'models/{model}')
+                m = gr.load(f'models/{model}')
             except Exception as error:
                 m = gr.Interface(lambda txt: None, ['text'], ['image'])
             models_load.update({model: m})
@@ -72,7 +67,10 @@ def img_to_img_fn(model_str, image, prompt, negative_prompt, max_retries=10):
     while retries < max_retries:
         try:
             noise = str(randint(0, 9999999))
-            result = models_load[model_str](prompt=f'{prompt} {noise}', negative_prompt=negative_prompt, image=image).images[0]
+            if hasattr(models_load[model_str], 'negative_prompt'):
+                result = models_load[model_str](image=image, prompt=f'{prompt} {noise}', negative_prompt=negative_prompt)
+            else:
+                result = models_load[model_str](image=image, prompt=f'{prompt} {noise}')
             return result
         except Exception as e:
             # Check for specific error messages or status codes