Update app.py

app.py

@@ -10,7 +10,7 @@ from huggingface_hub import InferenceClient
 
 # Project by Nymbo
 
-def query_with_auto_routing(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024):
+def query_with_auto_routing(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024, input_image=None):
     """
     Generate images using HF's automatic provider routing
     """
@@ -18,6 +18,7 @@ def query_with_auto_routing(prompt, model, custom_lora, is_negative=False, steps
     print(f"Prompt: {prompt}")
     print(f"Model: {model}")
     print(f"Custom LoRA: {custom_lora}")
+    print(f"Input Image: {input_image is not None}")
     print(f"Parameters - Steps: {steps}, CFG Scale: {cfg_scale}, Seed: {seed}, Strength: {strength}, Width: {width}, Height: {height}")
 
     # Check if the prompt is empty or None
@@ -54,15 +55,30 @@ def query_with_auto_routing(prompt, model, custom_lora, is_negative=False, steps
         enhanced_prompt = apply_model_prompt_enhancements(model, enhanced_prompt)
         
         # Generate image using automatic provider routing
-        image = client.text_to_image(
-            prompt=enhanced_prompt,
-            model=model_id,
-            width=width,
-            height=height,
-            num_inference_steps=steps,
-            guidance_scale=cfg_scale,
-            seed=seed if seed != -1 else None,
-        )
+        if input_image is not None:
+            # Image-to-image generation
+            print(f"Performing image-to-image generation with model: {model_id}")
+            image = client.image_to_image(
+                image=input_image,
+                prompt=enhanced_prompt,
+                model=model_id,
+                width=width,
+                height=height,
+                num_inference_steps=steps,
+                guidance_scale=cfg_scale,
+                # Note: strength parameter is handled differently in image_to_image
+            )
+        else:
+            # Text-to-image generation
+            image = client.text_to_image(
+                prompt=enhanced_prompt,
+                model=model_id,
+                width=width,
+                height=height,
+                num_inference_steps=steps,
+                guidance_scale=cfg_scale,
+                seed=seed if seed != -1 else None,
+            )
         
         print(f'Generation {key} completed with automatic routing!')
         return image
@@ -331,11 +347,11 @@ def apply_model_prompt_enhancements(model_name, prompt):
     
     return prompt
 
-def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024):
+def query(prompt, model, custom_lora, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024, input_image=None):
     """
     Main query function - now uses automatic provider routing
     """
-    return query_with_auto_routing(prompt, model, custom_lora, is_negative, steps, cfg_scale, sampler, seed, strength, width, height)
+    return query_with_auto_routing(prompt, model, custom_lora, is_negative, steps, cfg_scale, sampler, seed, strength, width, height, input_image)
 
 # Custom CSS to hide the footer in the interface
 css = """
@@ -354,12 +370,24 @@ with gr.Blocks(theme='Nymbo/Alyx_Theme') as dalle:
                 with gr.Row():
                     # Textbox for user to input the prompt
                     text_prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here", lines=3, elem_id="prompt-text-input")
+                with gr.Row():
+                    # Checkbox for enabling image editing
+                    enable_image_editing = gr.Checkbox(label="Enable Image Editing", value=False, elem_id="enable-image-editing")
                 with gr.Row():
                     # Textbox for custom LoRA input
                     custom_lora = gr.Textbox(label="Custom LoRA", info="LoRA Hugging Face path (optional)", placeholder="multimodalart/vintage-ads-flux")
+                with gr.Row():
+                    # Image upload component (hidden by default)
+                    image_upload = gr.Image(
+                        label="Upload Image for Editing",
+                        type="pil",
+                        visible=False,
+                        elem_id="image-upload"
+                    )
                 with gr.Row():
                     # Accordion for selecting the model
-                    with gr.Accordion("Featured Models", open=False):
+                    featured_models_accordion = gr.Accordion("Featured Models", open=False, elem_id="featured-models-accordion")
+                    with featured_models_accordion:
                         # Textbox for searching models
                         model_search = gr.Textbox(label="Filter Models", placeholder="Search for a featured model...", lines=1, elem_id="model-search-input")
                         models_list = (
@@ -481,6 +509,28 @@ with gr.Blocks(theme='Nymbo/Alyx_Theme') as dalle:
                         # Update model list when search box is used
                         model_search.change(filter_models, inputs=model_search, outputs=model)
 
+                # Function to toggle image editing mode
+                def toggle_image_editing(enable_editing):
+                    if enable_editing:
+                        return (
+                            gr.update(visible=True),  # Show image upload
+                            gr.update(visible=False),  # Hide featured models accordion
+                            gr.update(value="black-forest-labs/FLUX.1-Kontext-dev")  # Set custom LoRA
+                        )
+                    else:
+                        return (
+                            gr.update(visible=False),  # Hide image upload
+                            gr.update(visible=True),  # Show featured models accordion
+                            gr.update(value="")  # Clear custom LoRA
+                        )
+
+                # Set up the toggle functionality
+                enable_image_editing.change(
+                    toggle_image_editing,
+                    inputs=[enable_image_editing],
+                    outputs=[image_upload, featured_models_accordion, custom_lora]
+                )
+
     # Tab for advanced settings
     with gr.Tab("Advanced Settings"):
         with gr.Row():
@@ -601,7 +651,7 @@ with gr.Blocks(theme='Nymbo/Alyx_Theme') as dalle:
         image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")
 
     # Set up button click event to call the main query function
-    text_button.click(query, inputs=[text_prompt, model, custom_lora, negative_prompt, steps, cfg, method, seed, strength, width, height], outputs=image_output)
+    text_button.click(query, inputs=[text_prompt, model, custom_lora, negative_prompt, steps, cfg, method, seed, strength, width, height, image_upload], outputs=image_output)
 
 print("Launching Gradio interface...")  # Debug log
 # Launch the Gradio interface without showing the API or sharing externally