Update app.py

app.py

@@ -2,6 +2,8 @@ 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()
@@ -16,7 +18,10 @@ def load_fn(models):
     for model in models:
         if model not in models_load.keys():
             try:
-                m = gr.load(f'models/{model}')
+                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}')
             except Exception as error:
                 m = gr.Interface(lambda txt: None, ['text'], ['image'])
             models_load.update({model: m})
@@ -67,10 +72,7 @@ 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, prompt=f'{prompt} {noise}', negative_prompt=negative_prompt)
-            else:
-                result = models_load[model_str](image, prompt=f'{prompt} {noise}')
+            result = models_load[model_str](prompt=f'{prompt} {noise}', negative_prompt=negative_prompt, image=image).images[0]
             return result
         except Exception as e:
             # Check for specific error messages or status codes
@@ -148,6 +150,154 @@ def make_image_to_image():
 
 custom_css = """
 /* Your existing CSS styles here */
+:root {
+    --body-background-fill: #2d3d4f;
+}
+body {
+    background-color: var(--body-background-fill) !important;
+    color: #2d3d4f;
+    margin: 0;
+    padding: 0;
+    font-family: Arial, sans-serif;
+    height: 100vh;
+    overflow-y: auto;
+}
+.gradio-container {
+    background-color: #2d3d4f;
+    color: #c5c6c7;
+    padding: 20px;
+    border-radius: 8px;
+    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
+    width: 100%;
+    max-width: 1200px;
+    margin: 20px auto;
+    display: block;
+    min-height: 100vh;
+}
+.app_title {
+    background-color: #2d3d4f;
+    color: #c5c6c7;
+    padding: 10px 20px;
+    border-bottom: 1px solid #3b4252;
+    text-align: center;
+    font-size: 24px;
+    font-weight: bold;
+    width: 100%;
+    box-sizing: border-box;
+    margin-bottom: 20px;
+}
+.custom_textbox, .custom_negative_textbox {
+    background-color: #2d343f;
+    border: 1px solid #3b4252;
+    color: #7f8184;
+    padding: 10px;
+    border-radius: 4px;
+    margin-bottom: 10px;
+    width: 100%;
+    box-sizing: border-box;
+}
+.custom_gen_button {
+    background-color: #8b38ff;
+    border: 1px solid #ffffff;
+    color: blue;
+    padding: 15px 32px;
+    text-align: center;
+    text-decoration: none;
+    display: inline-block;
+    font-size: 16px;
+    margin: 4px 2px;
+    cursor: pointer;
+    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
+    transition: transform 0.2s, box-shadow 0.2s;
+    border-radius: 4px;
+}
+.custom_gen_button:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 6px 10px rgba(0, 0, 0, 0.3);
+}
+.custom_stop_button {
+    background-color: #6200ea;
+    border: 1px solid #ffffff;
+    color: blue;
+    padding: 15px 32px;
+    text-align: center;
+    text-decoration: none;
+    display: inline-block;
+    font-size: 16px;
+    margin: 4px 2px;
+    cursor: pointer;
+    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
+    transition: transform 0.2s, box-shadow 0.2s;
+    border-radius: 4px;
+}
+.custom_stop_button:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 6px 10px rgba(0, 0, 0, 0.3);
+}
+.custom_image {
+    border: 1px solid #3b4252;
+    background-color: #2d343f;
+    border-radius: 4px;
+    margin: 10px;
+    max-width: 100%;
+    box-sizing: border-box;
+}
+.custom_accordion {
+    background-color: #2d3d4f;
+    color: #7f8184;
+    border: 1px solid #3b4252;
+    border-radius: 4px;
+    margin-top: 20px;
+    width: 100%;
+    box-sizing: border-box;
+    transition: margin 0.2s ease;
+}
+.custom_accordion .gr-accordion-header {
+    background-color: #2d3d4f;
+    color: #7f8184;
+    padding: 10px 20px;
+    border-bottom: 1px solid #5b6270;
+    cursor: pointer;
+    font-size: 18px;
+    font-weight: bold;
+    height: 40px;
+    display: flex;
+    align-items: center;
+}
+.custom_accordion .gr-accordion-header:hover {
+    background-color: #2d3d4f;
+}
+.custom_accordion .gr-accordion-content {
+    padding: 10px 20px;
+    background-color: #2d3d4f;
+    border-top: 1px solid #5b6270;
+    max-height: 0;
+    overflow: hidden;
+    transition: max-height 0.2s ease;
+}
+.custom_accordion .gr-accordion-content.open {
+    max-height: 500px;
+}
+.custom_checkbox_group {
+    background-color: #2d343f;
+    border: 1px solid #3b4252;
+    color: #7f8184;
+    border-radius: 4px;
+    padding: 10px;
+    width: 100%;
+    box-sizing: border-box;
+}
+@media (max-width: 768px) {
+   .gradio-container {
+        width: 100%;
+        margin: 0;
+        padding: 10px;
+    }
+   .custom_textbox, .custom_negative_textbox, .custom_image, .custom_checkbox_group {
+        width: 100%;
+        box-sizing: border-box;
+    }
+}
 """
 
 with gr.Blocks(css=custom_css) as demo: