Update app.py

app.py

@@ -155,10 +155,7 @@ body {
 # --- Gradio Interface ---
 
 # Define demo OUTSIDE the blocks
-demo = gr.Blocks(css=css) 
-
-def start_generation():
-    return demo
+demo = gr.Blocks(css=css)
 
 # Initial Welcome Screen
 with gr.Blocks() as welcome_screen:
@@ -176,9 +173,6 @@ with gr.Blocks() as welcome_screen:
     with gr.Row():
         start_button = gr.Button("Get Started", variant='primary', elem_id="start-button")
 
-    # Connect Welcome Screen to Generation Screen 
-    start_button.click(fn=start_generation, outputs=welcome_screen) # Output welcome_screen
-
 # Main Generation Screen (Now use the defined 'demo')
 with demo:
     with gr.Row():
@@ -232,5 +226,16 @@ with demo:
 
     text_button.click(flip_text, inputs=[prompt, negative_prompt, model, steps, sampler, cfg_scale, width, height, seed], outputs=image_output)
 
+# Tabs to switch between screens
+tabs = gr.Tabs()
+with tabs:
+    with gr.TabItem("Welcome Screen"):
+        welcome_screen
+    with gr.TabItem("Main Generation Screen"):
+        demo
+
+# Connect Welcome Screen to Generation Screen 
+start_button.click(fn=lambda: tabs.set_selected("Main Generation Screen"))
+
 # Launch the Gradio app
-welcome_screen.launch(max_threads=128)
+tabs.launch(max_threads=128)