Update app.py

app.py

@@ -67,7 +67,6 @@ def image_to_base64(image_path):
         img_str = base64.b64encode(buffered.getvalue())
     return img_str.decode('utf-8')
 
-# Ensure that the API key is available
 api_key = os.getenv("PRODIA_API_KEY")
 if not api_key:
     raise ValueError("Prodia API key not found in environment variables")
@@ -160,6 +159,9 @@ body {
 with gr.Blocks(css=css) as demo:
     state = gr.State(value="Welcome Screen")
 
+    def update_visibility(tab_state):
+        return gr.Tabs.update(visible=(tab_state == "Main Generation Screen"))
+
     with gr.Tabs() as tabs:
         with gr.TabItem("Welcome Screen"):
             with gr.Row():
@@ -228,7 +230,7 @@ with gr.Blocks(css=css) as demo:
 
     text_button.click(flip_text, inputs=[prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed], outputs=image_output)
     start_button.click(fn=lambda: "Main Generation Screen", inputs=None, outputs=state)
-    state.change(fn=lambda x: gr.update(visible=(x == "Main Generation Screen")), inputs=state, outputs=tabs)
+    state.change(fn=update_visibility, inputs=state, outputs=tabs)
 
 # Launch the Gradio app
 demo.launch(max_threads=128)