Update app.py

app.py

@@ -59,7 +59,33 @@ def gen_fn(model_str, prompt, negative_prompt, max_retries=10):
     
     return None
 
-def make_me():
+def img_to_img_fn(model_str, image, prompt, negative_prompt, max_retries=10):
+    if model_str == 'NA' or image is None:
+        return None
+    
+    retries = 0
+    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}')
+            return result
+        except Exception as e:
+            # Check for specific error messages or status codes
+            if "CUDA out of memory" in str(e) or "500" in str(e):
+                print(f"CUDA out of memory or server error: {e}")
+            else:
+                print(f"Error generating image: {e}")
+            
+            retries += 1
+            if retries >= max_retries:
+                raise Exception(f"Failed to generate image after {max_retries} retries.")
+    
+    return None
+
+def make_text_to_image():
     with gr.Row():
         with gr.Column(scale=1):
             txt_input = gr.Textbox(label='Your prompt:', lines=3, container=False, elem_id="custom_textbox", placeholder="Prompt")
@@ -88,163 +114,48 @@ def make_me():
         model_choice = gr.CheckboxGroup(models, label=f'{num_models} different models selected', value=default_models, interactive=True, elem_id="custom_checkbox_group")
         model_choice.change(update_imgbox, model_choice, output)
         model_choice.change(extend_choices, model_choice, current_models)
+
+def make_image_to_image():
+    with gr.Row():
+        with gr.Column(scale=1):
+            img_input = gr.Image(label='Input Image', type='pil')
+            txt_input = gr.Textbox(label='Your prompt:', lines=3, container=False, elem_id="custom_textbox", placeholder="Prompt")
+            negative_txt_input = gr.Textbox(label='Negative prompt:', lines=3, container=False, elem_id="custom_negative_textbox", placeholder="Negative Prompt")
+            with gr.Row():
+                gen_button = gr.Button('Generate images', elem_id="custom_gen_button")
+                stop_button = gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
+                
+                def on_generate_click():
+                    return gr.Button('Generate images', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=True, elem_id="custom_stop_button")
+                
+                def on_stop_click():
+                    return gr.Button('Generate images', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
+                
+                gen_button.click(on_generate_click, inputs=None, outputs=[gen_button, stop_button])
+                stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button])
     
     with gr.Row():
-        gr.HTML("")
+        output = [gr.Image(label=m, min_width=250, height=250, elem_id="custom_image") for m in default_models]
+        current_models = [gr.Textbox(m, visible=False) for m in default_models]
+        for m, o in zip(current_models, output):
+            gen_event = gen_button.click(img_to_img_fn, [m, img_input, txt_input, negative_txt_input], o)
+            stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button], cancels=[gen_event])
+    
+    with gr.Accordion('Model selection', elem_id="custom_accordion"):
+        model_choice = gr.CheckboxGroup(models, label=f'{num_models} different models selected', value=default_models, interactive=True, elem_id="custom_checkbox_group")
+        model_choice.change(update_imgbox, model_choice, output)
+        model_choice.change(extend_choices, model_choice, current_models)
 
 custom_css = """
-: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;
-    }
-}
+/* Your existing CSS styles here */
 """
 
-with gr.Blocks(css=custom_css) as demo: 
-    make_me()
+with gr.Blocks(css=custom_css) as demo:
+    with gr.Tabs():
+        with gr.TabItem("Text-to-Image"):
+            make_text_to_image()
+        with gr.TabItem("Image-to-Image"):
+            make_image_to_image()
 
 demo.queue(concurrency_count=500)
 demo.launch()