Update app.py

app.py

@@ -2,6 +2,8 @@ import os
 import shlex
 import spaces
 import subprocess
+
+# --- Instalación del Entorno (sin cambios) ---
 def install_cuda_toolkit():
     CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.54.14_linux.run"
     CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
@@ -16,6 +18,7 @@ def install_cuda_toolkit():
         "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
     )
     os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
+
 install_cuda_toolkit()
 os.system("pip list | grep torch")
 os.system('nvcc -V')
@@ -23,6 +26,7 @@ print("cd /home/user/app/step1x3d_texture/differentiable_renderer/ && python set
 os.system("cd /home/user/app/step1x3d_texture/differentiable_renderer/ && python setup.py install")
 
 subprocess.run(shlex.split("pip install custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl"), check=True)
+
 import time
 import uuid
 import torch
@@ -36,7 +40,7 @@ from step1x3d_texture.pipelines.step1x_3d_texture_synthesis_pipeline import (
 )
 from step1x3d_geometry.models.pipelines.pipeline_utils import reduce_face, remove_degenerate_face
 
-
+# --- Configuración y Carga de Modelos (sin cambios) ---
 parser = argparse.ArgumentParser()
 parser.add_argument(
     "--geometry_model", type=str, default="Step1X-3D-Geometry-Label-1300m"
@@ -56,11 +60,16 @@ geometry_model = Step1X3DGeometryPipeline.from_pretrained(
 texture_model = Step1X3DTexturePipeline.from_pretrained("stepfun-ai/Step1X-3D", subfolder=args.texture_model)
 
 
-@spaces.GPU(duration=240)
-def generate_func(
+# --- NUEVAS FUNCIONES DE GENERACIÓN SEPARADAS ---
+
+@spaces.GPU(duration=180)
+def generate_geometry(
     input_image_path, guidance_scale, inference_steps, max_facenum, symmetry, edge_type
 ):
-    # geometry_model = geometry_model.to("cuda")
+    """
+    Función que genera únicamente la geometría.
+    """
+    print("Iniciando generación de geometría...")
     if "Label" in args.geometry_model:
         symmetry_values = ["x", "asymmetry"]
         out = geometry_model(
@@ -80,24 +89,52 @@ def generate_func(
         )
 
     save_name = str(uuid.uuid4())
-    print(save_name)
     geometry_save_path = f"{args.cache_dir}/{save_name}.glb"
     geometry_mesh = out.mesh[0]
     geometry_mesh.export(geometry_save_path)
-
+    
+    torch.cuda.empty_cache()
+    print(f"Geometría guardada en: {geometry_save_path}")
+    
+    # Devuelve la ruta para mostrar y para guardar en el estado
+    return geometry_save_path, geometry_save_path
+
+@spaces.GPU(duration=120)
+def generate_texture(input_image_path, geometry_path):
+    """
+    Función que aplica la textura a una geometría existente.
+    """
+    if not geometry_path or not os.path.exists(geometry_path):
+        raise gr.Error("Por favor, primero genera la geometría antes de texturizar.")
+        
+    print(f"Iniciando texturizado para la malla: {geometry_path}")
+    geometry_mesh = trimesh.load(geometry_path)
+    
+    # Post-procesamiento de la geometría antes de texturizar
     geometry_mesh = remove_degenerate_face(geometry_mesh)
     geometry_mesh = reduce_face(geometry_mesh)
+    
+    # Aplicar textura
     textured_mesh = texture_model(input_image_path, geometry_mesh)
+    
+    save_name = os.path.basename(geometry_path).replace(".glb", "")
     textured_save_path = f"{args.cache_dir}/{save_name}-textured.glb"
     textured_mesh.export(textured_save_path)
-
+    
     torch.cuda.empty_cache()
-    print("Generate finish")
-    return geometry_save_path, textured_save_path
+    print(f"Malla texturizada guardada en: {textured_save_path}")
+    
+    return textured_save_path
+
 
+# --- NUEVA INTERFAZ DE GRADIO ---
 
 with gr.Blocks(title="Step1X-3D demo") as demo:
     gr.Markdown("# Step1X-3D")
+    
+    # Estado para almacenar la ruta de la geometría generada
+    geometry_path_state = gr.State()
+
     with gr.Row():
         with gr.Column(scale=2):
             input_image = gr.Image(label="Image", type="filepath")
@@ -118,10 +155,15 @@ with gr.Blocks(title="Step1X-3D demo") as demo:
                 value="sharp",
                 type="value",
             )
-            btn = gr.Button("Start")
+            
+            with gr.Row():
+                btn_geo = gr.Button("1. Generate Geometry")
+                btn_tex = gr.Button("2. Generate Texture", visible=False)
+
         with gr.Column(scale=4):
-            textured_preview = gr.Model3D(label="Textured", height=380)
-            geometry_preview = gr.Model3D(label="Geometry", height=380)
+            textured_preview = gr.Model3D(label="Textured Model", height=380, clear_color=[0.0, 0.0, 0.0, 0.0])
+            geometry_preview = gr.Model3D(label="Geometry-only Model", height=380, clear_color=[0.0, 0.0, 0.0, 0.0])
+            
         with gr.Column(scale=1):
             gr.Examples(
                 examples=[
@@ -138,8 +180,11 @@ with gr.Blocks(title="Step1X-3D demo") as demo:
                 cache_examples=False,
             )
 
-    btn.click(
-        generate_func,
+    # Lógica de los botones
+    
+    # Al hacer clic en el botón de geometría
+    btn_geo.click(
+        fn=generate_geometry,
         inputs=[
             input_image,
             guidance_scale,
@@ -148,7 +193,21 @@ with gr.Blocks(title="Step1X-3D demo") as demo:
             symmetry,
             edge_type,
         ],
-        outputs=[geometry_preview, textured_preview],
+        outputs=[geometry_preview, geometry_path_state]
+    ).then(
+        # Cuando la geometría termine, actualiza la interfaz
+        fn=lambda: {
+            btn_tex: gr.update(visible=True), # Muestra el botón de texturizado
+            textured_preview: gr.update(value=None) # Limpia la vista previa de la textura anterior
+        },
+        outputs=[btn_tex, textured_preview]
+    )
+    
+    # Al hacer clic en el botón de textura
+    btn_tex.click(
+        fn=generate_texture,
+        inputs=[input_image, geometry_path_state],
+        outputs=[textured_preview],
     )
 
-demo.launch(ssr_mode=False)
+demo.launch(ssr_mode=False)