Update app.py

app.py

@@ -176,6 +176,7 @@ def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
     torch.cuda.empty_cache()
     return gaussian_path, gaussian_path
 
+# Gradio Interface
 # Gradio Interface
 with gr.Blocks() as demo:
     gr.Markdown("""
@@ -190,7 +191,6 @@ with gr.Blocks() as demo:
         with gr.Column():
             # Flux image generation inputs
             prompt = gr.Text(label="Prompt", placeholder="Enter your game asset description")
-            
             with gr.Accordion("Generation Settings", open=False):
                 seed = gr.Slider(0, MAX_SEED, label="Seed", value=42, step=1)
                 randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
@@ -199,64 +199,54 @@ with gr.Blocks() as demo:
                     height = gr.Slider(512, 1024, label="Height", value=1024, step=16)
                 with gr.Row():
                     guidance_scale = gr.Slider(0.0, 10.0, label="Guidance Scale", value=3.5, step=0.1)
-                #    num_inference_steps = gr.Slider(1, 50, label="Steps", value=8, step=1)
-
-            with gr.Accordion("3D Generation Settings", open=False):
-                gr.Markdown("Stage 1: Sparse Structure Generation")
-                with gr.Row():
-                    ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
-                    ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
-                gr.Markdown("Stage 2: Structured Latent Generation")
-                with gr.Row():
-                    slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
-                    slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
-
-            generate_btn = gr.Button("Generate")
-            
-            with gr.Accordion("GLB Extraction Settings", open=False):
-                mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01)
-                texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
             
-            with gr.Row():
-                extract_glb_btn = gr.Button("Extract GLB", interactive=False)
-                extract_gs_btn = gr.Button("Extract Gaussian", interactive=False)
+            # Botones separados
+            generate_image_btn = gr.Button("Generar Imagen")
+            generate_video_btn = gr.Button("Generar Video", interactive=False)
 
         with gr.Column():
             generated_image = gr.Image(label="Generated Asset", type="pil")
-            
-        with gr.Column():
-            
             video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True)
-            model_output = LitModel3D(label="Extracted GLB/Gaussian", exposure=8.0, height=400)
-            
-            with gr.Row():
-                download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
-                download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False)
-
+        
+        model_output = LitModel3D(label="Extracted GLB/Gaussian", exposure=8.0, height=400)
+    
+    with gr.Row():
+        extract_glb_btn = gr.Button("Extract GLB", interactive=False)
+        extract_gs_btn = gr.Button("Extract Gaussian", interactive=False)
+    
+    with gr.Row():
+        download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
+        download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False)
+    
+    # Estado para almacenar la imagen generada temporalmente
+    temp_image_state = gr.State()
     output_buf = gr.State()
 
     # Event handlers
     demo.load(start_session)
     demo.unload(end_session)
 
-    generate_btn.click(
+    # Generar imagen
+    generate_image_btn.click(
         generate_flux_image,
         inputs=[prompt, seed, randomize_seed, width, height, guidance_scale],
-        outputs=[generated_image],
+        outputs=[generated_image, temp_image_state],  # Almacenar la imagen en el estado temporal
     ).then(
+        lambda: gr.Button(interactive=True),  # Habilitar el botón "Generar Video"
+        outputs=[generate_video_btn],
+    )
+
+    # Generar video
+    generate_video_btn.click(
         image_to_3d,
-        inputs=[generated_image, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
+        inputs=[temp_image_state, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
         outputs=[output_buf, video_output],
     ).then(
         lambda: tuple([gr.Button(interactive=True), gr.Button(interactive=True)]),
         outputs=[extract_glb_btn, extract_gs_btn],
     )
 
-    video_output.clear(
-        lambda: tuple([gr.Button(interactive=False), gr.Button(interactive=False)]),
-        outputs=[extract_glb_btn, extract_gs_btn],
-    )
-
+    # Extraer GLB
     extract_glb_btn.click(
         extract_glb,
         inputs=[output_buf, mesh_simplify, texture_size],
@@ -265,7 +255,8 @@ with gr.Blocks() as demo:
         lambda: gr.Button(interactive=True),
         outputs=[download_glb],
     )
-    
+
+    # Extraer Gaussian
     extract_gs_btn.click(
         extract_gaussian,
         inputs=[output_buf],
@@ -275,11 +266,6 @@ with gr.Blocks() as demo:
         outputs=[download_gs],
     )
 
-    model_output.clear(
-        lambda: gr.Button(interactive=False),
-        outputs=[download_glb],
-    )
-
 # Initialize both pipelines
 if __name__ == "__main__":
     from diffusers import FluxTransformer2DModel, FluxPipeline, BitsAndBytesConfig, GGUFQuantizationConfig