Refactor flux_events.py to add GPU duration for image generation

modules/events/flux_events.py

@@ -206,7 +206,7 @@ def remove_from_enabled_loras(enabled_loras, index):
     )
 
 
-@spaces.GPU
+@spaces.GPU(duration=120)
 def generate_image(
         model, prompt, fast_generation, enabled_loras,
         lora_slider_0, lora_slider_1, lora_slider_2, lora_slider_3, lora_slider_4, lora_slider_5,

tabs/image_tab.py

@@ -11,136 +11,137 @@ from config import flux_models, sdxl_models, flux_loras
 
 def image_tab():
     with gr.Tabs():
-        flux_tab()
-        sdxl_tab()
+        with gr.Tab("Flux"):
+            flux_tab()
+        with gr.Tab("SDXL"):
+            sdxl_tab()
 
 
 def flux_tab():
     loras = flux_loras
-    with gr.Tab("Flux"):
-        with gr.Row():
-            with gr.Column():
-                with gr.Group() as image_options:
-                    model = gr.Dropdown(label="Models", choices=flux_models, value=flux_models[0], interactive=True)
-                    prompt = gr.Textbox(lines=5, label="Prompt")
-                    fast_generation = gr.Checkbox(label="Fast Generation (Hyper-SD) 🧪")
-                
+    with gr.Row():
+        with gr.Column():
+            with gr.Group() as image_options:
+                model = gr.Dropdown(label="Models", choices=flux_models, value=flux_models[0], interactive=True)
+                prompt = gr.Textbox(lines=5, label="Prompt")
+                fast_generation = gr.Checkbox(label="Fast Generation (Hyper-SD) 🧪")
+            
+            
+            with gr.Accordion("Loras", open=True): # Lora Gallery
+                lora_gallery = gr.Gallery(
+                    label="Gallery",
+                    value=[(lora['image'], lora['title']) for lora in loras],
+                    allow_preview=False,
+                    columns=3,
+                    rows=3,
+                    type="pil"
+                )
                 
-                with gr.Accordion("Loras", open=True): # Lora Gallery
-                    lora_gallery = gr.Gallery(
-                        label="Gallery",
-                        value=[(lora['image'], lora['title']) for lora in loras],
-                        allow_preview=False,
-                        columns=3,
-                        rows=3,
-                        type="pil"
-                    )
-                    
-                    with gr.Group():
-                        with gr.Column():
+                with gr.Group():
+                    with gr.Column():
+                        with gr.Row():
+                            custom_lora = gr.Textbox(label="Custom Lora", info="Enter a Huggingface repo path")
+                            selected_lora = gr.Textbox(label="Selected Lora", info="Choose from the gallery or enter a custom LoRA")
+                        
+                        custom_lora_info = gr.HTML(visible=False)
+                        add_lora = gr.Button(value="Add LoRA")
+                        
+                        enabled_loras = gr.State(value=[])
+                        with gr.Group():
                             with gr.Row():
-                                custom_lora = gr.Textbox(label="Custom Lora", info="Enter a Huggingface repo path")
-                                selected_lora = gr.Textbox(label="Selected Lora", info="Choose from the gallery or enter a custom LoRA")
-                            
-                            custom_lora_info = gr.HTML(visible=False)
-                            add_lora = gr.Button(value="Add LoRA")
-                            
-                            enabled_loras = gr.State(value=[])
-                            with gr.Group():
-                                with gr.Row():
-                                    for i in range(6): # only support max 6 loras due to inference time
+                                for i in range(6): # only support max 6 loras due to inference time
+                                    with gr.Column():
+                                        with gr.Column(scale=2):
+                                            globals()[f"lora_slider_{i}"] = gr.Slider(label=f"LoRA {i+1}", minimum=0, maximum=1, step=0.01, value=0.8, visible=False, interactive=True)
                                         with gr.Column():
-                                            with gr.Column(scale=2):
-                                                globals()[f"lora_slider_{i}"] = gr.Slider(label=f"LoRA {i+1}", minimum=0, maximum=1, step=0.01, value=0.8, visible=False, interactive=True)
-                                            with gr.Column():
-                                                globals()[f"lora_remove_{i}"] = gr.Button(value="Remove LoRA", visible=False)
+                                            globals()[f"lora_remove_{i}"] = gr.Button(value="Remove LoRA", visible=False)
 
-                
-                with gr.Accordion("Embeddings", open=False): # Embeddings
-                    gr.Label("To be implemented")
-                
-                
-                with gr.Accordion("Image Options"): # Image Options
-                    with gr.Tabs():
-                        image_options = {
-                            "img2img": "Upload Image",
-                            "inpaint": "Upload Image",
-                            "canny": "Upload Image",
-                            "pose": "Upload Image",
-                            "depth": "Upload Image",
-                        }
-                        
-                        for image_option, label in image_options.items():
-                            with gr.Tab(image_option):
-                                if not image_option in ['inpaint', 'scribble']:
-                                    globals()[f"{image_option}_image"] = gr.Image(label=label, type="pil")
-                                elif image_option in ['inpaint', 'scribble']:
-                                    globals()[f"{image_option}_image"] = gr.ImageEditor(
-                                        label=label,
-                                        image_mode='RGB',
-                                        layers=False,
-                                        brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed") if image_option == 'inpaint' else gr.Brush(),
-                                        interactive=True,
-                                        type="pil",
-                                    )
-                                
-                                # Image Strength (Co-relates to controlnet strength, strength for img2img n inpaint)
-                                globals()[f"{image_option}_strength"] = gr.Slider(label="Strength", minimum=0, maximum=1, step=0.01, value=1.0, interactive=True)
-                        
-                        resize_mode = gr.Radio(
-                            label="Resize Mode",
-                            choices=["crop and resize", "resize only", "resize and fill"],
-                            value="resize and fill",
-                            interactive=True
-                        )
             
+            with gr.Accordion("Embeddings", open=False): # Embeddings
+                gr.Label("To be implemented")
             
-            with gr.Column():
-                with gr.Group():
-                    output_images = gr.Gallery(
-                            label="Output Images",
-                            value=[],
-                            allow_preview=True,
-                            type="pil",
-                            interactive=False,
-                        )
-                    generate_images = gr.Button(value="Generate Images", variant="primary")            
-                
-                with gr.Accordion("Advance Settings", open=True):
-                    with gr.Row():
-                        scheduler = gr.Dropdown(
-                            label="Scheduler",
-                            choices = [
-                                "fm_euler"
-                            ],
-                            value="fm_euler",
-                            interactive=True
-                        )
-
-                    with gr.Row():
-                        for column in range(2):
-                            with gr.Column():
-                                options = [
-                                    ("Height", "image_height", 64, 1024, 64, 1024, True),
-                                    ("Width", "image_width", 64, 1024, 64, 1024, True),
-                                    ("Num Images Per Prompt", "image_num_images_per_prompt", 1, 4, 1, 1, True),
-                                    ("Num Inference Steps", "image_num_inference_steps", 1, 100, 1, 20, True),
-                                    ("Clip Skip", "image_clip_skip", 0, 2, 1, 2, False),
-                                    ("Guidance Scale", "image_guidance_scale", 0, 20, 0.5, 3.5, True),
-                                    ("Seed", "image_seed", 0, 100000, 1, random.randint(0, 100000), True),
-                                ]
-                                for label, var_name, min_val, max_val, step, value, visible in options[column::2]:
-                                    globals()[var_name] = gr.Slider(label=label, minimum=min_val, maximum=max_val, step=step, value=value, visible=visible, interactive=True)
+            
+            with gr.Accordion("Image Options", open=False): # Image Options
+                with gr.Tabs():
+                    image_options = {
+                        "img2img": "Upload Image",
+                        "inpaint": "Upload Image",
+                        "canny": "Upload Image",
+                        "pose": "Upload Image",
+                        "depth": "Upload Image",
+                    }
                     
-                    with gr.Row():
-                        refiner = gr.Checkbox(
-                            label="Refiner 🧪",
-                            value=False,
-                        )
-                        vae = gr.Checkbox(
-                            label="VAE",
-                            value=True,
-                        )
+                    for image_option, label in image_options.items():
+                        with gr.Tab(image_option):
+                            if not image_option in ['inpaint', 'scribble']:
+                                globals()[f"{image_option}_image"] = gr.Image(label=label, type="pil")
+                            elif image_option in ['inpaint', 'scribble']:
+                                globals()[f"{image_option}_image"] = gr.ImageEditor(
+                                    label=label,
+                                    image_mode='RGB',
+                                    layers=False,
+                                    brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed") if image_option == 'inpaint' else gr.Brush(),
+                                    interactive=True,
+                                    type="pil",
+                                )
+                            
+                            # Image Strength (Co-relates to controlnet strength, strength for img2img n inpaint)
+                            globals()[f"{image_option}_strength"] = gr.Slider(label="Strength", minimum=0, maximum=1, step=0.01, value=1.0, interactive=True)
+                    
+                    resize_mode = gr.Radio(
+                        label="Resize Mode",
+                        choices=["crop and resize", "resize only", "resize and fill"],
+                        value="resize and fill",
+                        interactive=True
+                    )
+        
+        
+        with gr.Column():
+            with gr.Group():
+                output_images = gr.Gallery(
+                        label="Output Images",
+                        value=[],
+                        allow_preview=True,
+                        type="pil",
+                        interactive=False,
+                    )
+                generate_images = gr.Button(value="Generate Images", variant="primary")            
+            
+            with gr.Accordion("Advance Settings", open=True):
+                with gr.Row():
+                    scheduler = gr.Dropdown(
+                        label="Scheduler",
+                        choices = [
+                            "fm_euler"
+                        ],
+                        value="fm_euler",
+                        interactive=True
+                    )
+
+                with gr.Row():
+                    for column in range(2):
+                        with gr.Column():
+                            options = [
+                                ("Height", "image_height", 64, 1024, 64, 1024, True),
+                                ("Width", "image_width", 64, 1024, 64, 1024, True),
+                                ("Num Images Per Prompt", "image_num_images_per_prompt", 1, 4, 1, 1, True),
+                                ("Num Inference Steps", "image_num_inference_steps", 1, 100, 1, 20, True),
+                                ("Clip Skip", "image_clip_skip", 0, 2, 1, 2, False),
+                                ("Guidance Scale", "image_guidance_scale", 0, 20, 0.5, 3.5, True),
+                                ("Seed", "image_seed", 0, 100000, 1, random.randint(0, 100000), True),
+                            ]
+                            for label, var_name, min_val, max_val, step, value, visible in options[column::2]:
+                                globals()[var_name] = gr.Slider(label=label, minimum=min_val, maximum=max_val, step=step, value=value, visible=visible, interactive=True)
+                
+                with gr.Row():
+                    refiner = gr.Checkbox(
+                        label="Refiner 🧪",
+                        value=False,
+                    )
+                    vae = gr.Checkbox(
+                        label="VAE",
+                        value=True,
+                    )
     
     # Events
     # Base Options
@@ -179,5 +180,4 @@ def flux_tab():
 
 
 def sdxl_tab():
-    with gr.Tab("SDXL"):
-        gr.Label("To be implemented")
+    gr.Label("To be implemented")