Update app.py

app.py

@@ -3,12 +3,21 @@ import shlex
 import spaces
 import subprocess
 
-# --- Instalación del Entorno (sin cambios) ---
+# --------------------------------------------------------------------------
+# 1. INSTALACIÓN DEL ENTORNO Y DEPENDENCIAS
+# Esta sección es crucial para que el Space de Hugging Face funcione.
+# Instala el toolkit de CUDA y compila las extensiones C++/CUDA necesarias.
+# --------------------------------------------------------------------------
 def install_cuda_toolkit():
+    """Instala el toolkit de CUDA necesario para compilar las extensiones."""
     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)
+    
+    print("Descargando CUDA Toolkit...")
     subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
     subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
+    
+    print("Instalando CUDA Toolkit...")
     subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
 
     os.environ["CUDA_HOME"] = "/usr/local/cuda"
@@ -18,16 +27,21 @@ 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"
+    print("Configuración de entorno de CUDA finalizada.")
 
 install_cuda_toolkit()
+
+print("Verificando instalación de PyTorch y versión de NVCC:")
 os.system("pip list | grep torch")
 os.system('nvcc -V')
-print("cd /home/user/app/step1x3d_texture/differentiable_renderer/ && python setup.py install")
+
+print("Compilando extensión de renderizador diferenciable...")
 os.system("cd /home/user/app/step1x3d_texture/differentiable_renderer/ && python setup.py install")
 
+print("Instalando rasterizador personalizado...")
 subprocess.run(shlex.split("pip install custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl"), check=True)
+print("Instalación y compilación completadas.")
 
-import time
 import uuid
 import torch
 import trimesh
@@ -40,7 +54,10 @@ 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) ---
+# --------------------------------------------------------------------------
+# 2. CONFIGURACIÓN Y CARGA DE MODELOS
+# Aquí se definen los modelos a utilizar y se cargan en memoria.
+# --------------------------------------------------------------------------
 parser = argparse.ArgumentParser()
 parser.add_argument(
     "--geometry_model", type=str, default="Step1X-3D-Geometry-Label-1300m"
@@ -52,24 +69,37 @@ parser.add_argument("--cache_dir", type=str, default="cache")
 args = parser.parse_args()
 
 os.makedirs(args.cache_dir, exist_ok=True)
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
+print(f"Cargando modelo de geometría: {args.geometry_model}...")
 geometry_model = Step1X3DGeometryPipeline.from_pretrained(
     "stepfun-ai/Step1X-3D", subfolder=args.geometry_model
-).to("cuda")
+).to(device)
+print("Modelo de geometría cargado.")
 
+print(f"Cargando modelo de textura: {args.texture_model}...")
 texture_model = Step1X3DTexturePipeline.from_pretrained("stepfun-ai/Step1X-3D", subfolder=args.texture_model)
+print("Modelo de textura cargado.")
 
 
-# --- NUEVAS FUNCIONES DE GENERACIÓN SEPARADAS ---
+# --------------------------------------------------------------------------
+# 3. FUNCIONES DE GENERACIÓN SEPARADAS
+# Se divide la lógica en dos funciones: una para la geometría y otra para la textura.
+# --------------------------------------------------------------------------
 
 @spaces.GPU(duration=180)
 def generate_geometry(
-    input_image_path, guidance_scale, inference_steps, max_facenum, symmetry, edge_type
+    input_image_path, guidance_scale, inference_steps, max_facenum, symmetry, edge_type, progress=gr.Progress(track_tqdm=True)
 ):
     """
-    Función que genera únicamente la geometría.
+    Función que genera únicamente la geometría a partir de la imagen de entrada.
     """
+    if input_image_path is None:
+        raise gr.Error("Por favor, sube una imagen para empezar.")
+        
     print("Iniciando generación de geometría...")
+    
+    # Elige el pipeline adecuado basado en el nombre del modelo
     if "Label" in args.geometry_model:
         symmetry_values = ["x", "asymmetry"]
         out = geometry_model(
@@ -88,6 +118,7 @@ def generate_geometry(
             max_facenum=int(max_facenum),
         )
 
+    # Guarda el resultado en un fichero temporal
     save_name = str(uuid.uuid4())
     geometry_save_path = f"{args.cache_dir}/{save_name}.glb"
     geometry_mesh = out.mesh[0]
@@ -96,13 +127,13 @@ def generate_geometry(
     torch.cuda.empty_cache()
     print(f"Geometría guardada en: {geometry_save_path}")
     
-    # Devuelve la ruta para mostrar y para guardar en el estado
+    # Devuelve la ruta para mostrarla en el visor y para guardarla en el estado
     return geometry_save_path, geometry_save_path
 
 @spaces.GPU(duration=120)
-def generate_texture(input_image_path, geometry_path):
+def generate_texture(input_image_path, geometry_path, progress=gr.Progress(track_tqdm=True)):
     """
-    Función que aplica la textura a una geometría existente.
+    Función que aplica la textura a una geometría ya generada.
     """
     if not geometry_path or not os.path.exists(geometry_path):
         raise gr.Error("Por favor, primero genera la geometría antes de texturizar.")
@@ -110,13 +141,14 @@ def generate_texture(input_image_path, geometry_path):
     print(f"Iniciando texturizado para la malla: {geometry_path}")
     geometry_mesh = trimesh.load(geometry_path)
     
-    # Post-procesamiento de la geometría antes de texturizar
+    # Post-procesamiento opcional de la malla antes de texturizar
     geometry_mesh = remove_degenerate_face(geometry_mesh)
     geometry_mesh = reduce_face(geometry_mesh)
     
-    # Aplicar textura
+    # Llama al pipeline de texturizado
     textured_mesh = texture_model(input_image_path, geometry_mesh)
     
+    # Guarda el resultado final
     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)
@@ -127,42 +159,48 @@ def generate_texture(input_image_path, geometry_path):
     return textured_save_path
 
 
-# --- NUEVA INTERFAZ DE GRADIO ---
+# --------------------------------------------------------------------------
+# 4. INTERFAZ DE USUARIO CON GRADIO
+# Se define la apariencia y el comportamiento de la aplicación web.
+# --------------------------------------------------------------------------
 
 with gr.Blocks(title="Step1X-3D demo") as demo:
     gr.Markdown("# Step1X-3D")
+    gr.Markdown("### Demostración de generación de modelos 3D a partir de una única imagen")
     
-    # Estado para almacenar la ruta de la geometría generada
+    # Componente de estado: guarda la ruta de la geometría entre los pasos
     geometry_path_state = gr.State()
 
     with gr.Row():
         with gr.Column(scale=2):
-            input_image = gr.Image(label="Image", type="filepath")
-            guidance_scale = gr.Number(label="Guidance Scale", value="7.5")
-            inference_steps = gr.Slider(
-                label="Inferece Steps", minimum=1, maximum=100, value=50
-            )
-            max_facenum = gr.Number(label="Max Face Num", value="400000")
-            symmetry = gr.Radio(
-                choices=["symmetry", "asymmetry"],
-                label="Symmetry Type",
-                value="symmetry",
-                type="index",
-            )
-            edge_type = gr.Radio(
-                choices=["sharp", "normal", "smooth"],
-                label="Edge Type",
-                value="sharp",
-                type="value",
-            )
+            input_image = gr.Image(label="Imagen de Entrada", type="filepath")
+            
+            with gr.Accordion(label="Parámetros de Generación", open=True):
+                guidance_scale = gr.Number(label="Guidance Scale", value="7.5")
+                inference_steps = gr.Slider(
+                    label="Pasos de Inferencia", minimum=1, maximum=100, value=50, step=1
+                )
+                max_facenum = gr.Number(label="Máx. Número de Caras", value="400000")
+                symmetry = gr.Radio(
+                    choices=["symmetry", "asymmetry"],
+                    label="Tipo de Simetría",
+                    value="symmetry",
+                    type="index",
+                )
+                edge_type = gr.Radio(
+                    choices=["sharp", "normal", "smooth"],
+                    label="Tipo de Borde",
+                    value="sharp",
+                    type="value",
+                )
             
             with gr.Row():
-                btn_geo = gr.Button("1. Generate Geometry")
-                btn_tex = gr.Button("2. Generate Texture", visible=False)
+                btn_geo = gr.Button("1. Generar Geometría", variant="primary")
+                btn_tex = gr.Button("2. Generar Textura", visible=False, variant="primary")
 
         with gr.Column(scale=4):
-            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])
+            textured_preview = gr.Model3D(label="Modelo con Textura", height=380, clear_color=[0.0, 0.0, 0.0, 0.0])
+            geometry_preview = gr.Model3D(label="Modelo (solo geometría)", height=380, clear_color=[0.0, 0.0, 0.0, 0.0])
             
         with gr.Column(scale=1):
             gr.Examples(
@@ -180,9 +218,9 @@ with gr.Blocks(title="Step1X-3D demo") as demo:
                 cache_examples=False,
             )
 
-    # Lógica de los botones
+    # --- Lógica de los botones y el flujo de la interfaz ---
     
-    # Al hacer clic en el botón de geometría
+    # 1. Cuando el usuario hace clic en "Generar Geometría"
     btn_geo.click(
         fn=generate_geometry,
         inputs=[
@@ -195,19 +233,20 @@ with gr.Blocks(title="Step1X-3D demo") as demo:
         ],
         outputs=[geometry_preview, geometry_path_state]
     ).then(
-        # Cuando la geometría termine, actualiza la interfaz
+        # 2. Cuando la geometría termina, se ejecuta esta parte
         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
+            btn_tex: gr.update(visible=True),      # Hace visible el botón de texturizado
+            textured_preview: gr.update(value=None) # Limpia la vista previa de texturas anterior
         },
         outputs=[btn_tex, textured_preview]
     )
     
-    # Al hacer clic en el botón de textura
+    # 3. Cuando el usuario hace clic en "Generar Textura"
     btn_tex.click(
         fn=generate_texture,
         inputs=[input_image, geometry_path_state],
         outputs=[textured_preview],
     )
 
+# Lanza la aplicación
 demo.launch(ssr_mode=False)