Update app.py

app.py

@@ -1,15 +1,12 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 
-
-
 def update_model(selected_model, use_textbox, custom_model):
     if use_textbox:
         return custom_model  # Return the custom model from textbox
     else:
         return selected_model  # Return the selected model from dropdown
 
-
 # Initialize Hugging Face Inference Client
 def get_client(model_name):
     return InferenceClient(model=model_name)
@@ -32,13 +29,14 @@ with gr.Blocks() as demo:
     }
 
     gr.Markdown("# Text to Image Generator using Hugging Face Inference Client")
+    
     with gr.Row():
         use_textbox = gr.Checkbox(label="Use Custom Model", value=False)
         
     with gr.Row():
         # Input for text prompt
         prompt_input = gr.Textbox(label="Enter your prompt", placeholder="Describe the image you want...")
-            
+        
         # Button to generate image
         generate_button = gr.Button("Generate Image")
         
@@ -49,20 +47,19 @@ with gr.Blocks() as demo:
                 choices=list(model_options.keys()),  # Display model names
                 value="FLUX 1.0 (black-forest-labs)",  # Default model
             )
-        custom_model_textbox = gr.Textbox(label="Enter Custom Model", placeholder="Type your model name here")
+        
+        custom_model_textbox = gr.Textbox(label="Enter Custom Model", placeholder="Type your model name here", visible=False)
 
         output = gr.Textbox(label="Selected Model", interactive=False)
 
-
-
         with gr.Row():
             submit_button = gr.Button("Submit Option")
+        
         with gr.Column():
             # Image output
             image_output = gr.Image(label="Generated Image")
 
-
-   # Event to update the selected model based on the checkbox
+    # Event to update the selected model based on the checkbox
     use_textbox.change(
         lambda checked: model_dropdown.update(visible=not checked, value=None),
         inputs=[use_textbox],
@@ -76,15 +73,18 @@ with gr.Blocks() as demo:
     )
     
     # Button to confirm selection
-    
     submit_button.click(
         update_model,
         inputs=[model_dropdown, use_textbox, custom_model_textbox],
         outputs=[output]
-    ) 
+    )
     
-    # Link the button click to the function
-    generate_button.click(generate_image, inputs=[prompt_input, model_dropdown], outputs=image_output)
+    # Link the button click to the image generation function
+    generate_button.click(
+        generate_image,
+        inputs=[prompt_input, output],
+        outputs=image_output
+    )
 
 # Launch the Gradio app
 demo.launch()