Update app.py

app.py

@@ -2,57 +2,76 @@ import os
 import shlex
 import spaces
 import subprocess
-import uuid
-import torch
-import trimesh
 import argparse
-import numpy as np
-import gradio as gr
+import uuid
 import random
+import numpy as np
+import shutil
 
 # --------------------------------------------------------------------------
-# 1. INSTALACIÓN DEL ENTORNO Y DEPENDENCIAS (Sin cambios)
+# BLOQUE 1: INSTALACIÓN DEL ENTORNO Y DEPENDENCIAS
+# Esta sección DEBE ejecutarse en su totalidad ANTES de importar torch o gradio.
 # --------------------------------------------------------------------------
+
 @spaces.GPU(duration=120)
 def setup_environment():
-    """Prepara el entorno, instalando CUDA y compilando extensiones."""
+    """
+    Prepara todo el entorno necesario, instalando CUDA y compilando las extensiones.
+    Esta función se ejecuta una vez al inicio del Space.
+    """
     print("--- INICIANDO CONFIGURACIÓN DEL ENTORNO ---")
+
+    # 1. Instalar CUDA Toolkit
     if not os.path.exists("/usr/local/cuda"):
         print("Instalando 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 = f"/tmp/{os.path.basename(CUDA_TOOLKIT_URL)}"
+        
         subprocess.run(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE], check=True)
         subprocess.run(["chmod", "+x", CUDA_TOOLKIT_FILE], check=True)
         subprocess.run([CUDA_TOOLKIT_FILE, "--silent", "--toolkit", "--override"], check=True)
+        print("CUDA Toolkit instalado.")
     else:
         print("CUDA Toolkit ya está instalado.")
-    
+
+    # 2. Configurar variables de entorno
     os.environ["CUDA_HOME"] = "/usr/local/cuda"
     os.environ["PATH"] = f"{os.environ['CUDA_HOME']}/bin:{os.environ['PATH']}"
     os.environ["LD_LIBRARY_PATH"] = f"{os.environ['CUDA_HOME']}/lib64:{os.environ.get('LD_LIBRARY_PATH', '')}"
     os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
     print("Variables de entorno de CUDA configuradas.")
+
+    # 3. Verificar NVCC
+    print("Verificando versión de NVCC:")
     os.system('nvcc -V')
 
+    # 4. Compilar extensiones C++/CUDA
     print("Compilando extensión de renderizador diferenciable...")
     try:
-        subprocess.run("cd /home/user/app/step1x3d_texture/differentiable_renderer/ && python setup.py install", shell=True, check=True)
+        subprocess.run(
+            "cd /home/user/app/step1x3d_texture/differentiable_renderer/ && python setup.py install", 
+            shell=True, check=True, capture_output=True, text=True
+        )
+        print("Renderizador diferenciable compilado.")
     except subprocess.CalledProcessError as e:
-        print(f"ERROR al compilar el renderizador diferenciable: {e}")
+        print(f"ERROR al compilar el renderizador diferenciable:\n{e.stderr}")
         raise RuntimeError("Fallo crítico en la compilación del renderizador.")
 
     print("Instalando rasterizador personalizado...")
     try:
         subprocess.run(shlex.split("pip install custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl"), check=True)
+        print("Rasterizador personalizado instalado.")
     except Exception as e:
         print(f"ERROR al instalar el rasterizador personalizado: {e}")
         raise RuntimeError("Fallo crítico en la instalación del rasterizador.")
+
     print("--- CONFIGURACIÓN DEL ENTORNO FINALIZADA ---")
 
+# Ejecutar la configuración del entorno antes de cualquier otra cosa
 setup_environment()
 
 # --------------------------------------------------------------------------
-# 2. IMPORTACIONES Y CARGA DE MODELOS (Después de la configuración)
+# BLOQUE 2: IMPORTACIONES DE LA APLICACIÓN Y LÓGICA PRINCIPAL
 # --------------------------------------------------------------------------
 import torch
 import trimesh
@@ -60,15 +79,20 @@ import gradio as gr
 from step1x3d_texture.pipelines.step1x_3d_texture_synthesis_pipeline import Step1X3DTexturePipeline
 from step1x3d_geometry.models.pipelines.pipeline_utils import reduce_face, remove_degenerate_face
 
+# --------------------------------------------------------------------------
+# 3. CONFIGURACIÓN Y CARGA DEL MODELO DE TEXTURA
+# --------------------------------------------------------------------------
 parser = argparse.ArgumentParser()
 parser.add_argument("--texture_model", type=str, default="Step1X-3D-Texture")
 parser.add_argument("--cache_dir", type=str, default="cache")
 args = parser.parse_args()
 
-os.makedirs(args.cache_dir, exist_ok=True)
+TMP_DIR = args.cache_dir
+os.makedirs(TMP_DIR, exist_ok=True)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_SEED = np.iinfo(np.int32).max
 
+print(f"Dispositivo detectado: {device}")
 if not torch.cuda.is_available():
     raise RuntimeError("CUDA no está disponible para PyTorch. La aplicación no puede continuar.")
 
@@ -77,59 +101,82 @@ texture_model = Step1X3DTexturePipeline.from_pretrained("stepfun-ai/Step1X-3D",
 print("Modelo de textura cargado y listo.")
 
 # --------------------------------------------------------------------------
-# 3. FUNCIONES DE LÓGICA
+# 4. GESTIÓN DE SESIONES Y ARCHIVOS TEMPORALES
+# --------------------------------------------------------------------------
+
+def start_session(req: gr.Request):
+    """Crea un directorio temporal único para la sesión del usuario."""
+    session_hash = str(req.session_hash)
+    user_dir = os.path.join(TMP_DIR, session_hash)
+    os.makedirs(user_dir, exist_ok=True)
+    print(f"Sesión iniciada: {session_hash}. Directorio temporal: {user_dir}")
+
+def end_session(req: gr.Request):
+    """Elimina el directorio temporal de la sesión para liberar espacio."""
+    session_hash = str(req.session_hash)
+    user_dir = os.path.join(TMP_DIR, session_hash)
+    if os.path.exists(user_dir):
+        try:
+            shutil.rmtree(user_dir)
+            print(f"Sesión finalizada: {session_hash}. Directorio temporal eliminado.")
+        except Exception as e:
+            print(f"Error al eliminar el directorio de la sesión {session_hash}: {e}")
+
+# --------------------------------------------------------------------------
+# 5. FUNCIÓN DE GENERACIÓN DE TEXTURA
 # --------------------------------------------------------------------------
 
 def get_random_seed(randomize_seed, seed):
-    """Devuelve una semilla aleatoria si la casilla está marcada, de lo contrario devuelve la semilla del slider."""
+    """Genera una semilla aleatoria si se solicita."""
     if randomize_seed:
-        new_seed = random.randint(0, MAX_SEED)
-        print(f"Generando semilla aleatoria: {new_seed}")
-        return new_seed
-    else:
-        print(f"Usando semilla fija: {int(seed)}")
-        return int(seed)
+        return random.randint(0, MAX_SEED)
+    return int(seed)
 
 @spaces.GPU(duration=180)
-def generate_texture(
+def generate_texture_for_user_mesh(
     input_image_path,
     input_mesh_path,
     guidance_scale,
     inference_steps,
-    seed, # Ahora este valor es el definitivo (aleatorio o fijo)
+    seed,
+    randomize_seed,
     reference_conditioning_scale,
+    req: gr.Request, # Se añade el objeto request para obtener el hash de la sesión
     progress=gr.Progress(track_tqdm=True)
 ):
     """
-    Función principal que genera la textura. Ya no necesita el parámetro 'randomize_seed'.
+    Función principal que genera la textura para un modelo 3D proporcionado por el usuario.
     """
     if input_image_path is None:
         raise gr.Error("Por favor, sube una imagen de referencia para empezar.")
     if input_mesh_path is None:
         raise gr.Error("Por favor, sube un modelo 3D (.glb o .obj) para texturizar.")
 
-    print("Iniciando generación de textura para el modelo del usuario...")
+    session_hash = str(req.session_hash)
+    user_dir = os.path.join(TMP_DIR, session_hash)
+    print(f"[{session_hash}] Iniciando generación de textura...")
+    
+    # Obtener semilla
+    final_seed = get_random_seed(randomize_seed, seed)
     
-    # Actualizar la configuración del pipeline con los valores de la UI
+    # Actualizar la configuración del pipeline
     texture_model.config.guidance_scale = float(guidance_scale)
     texture_model.config.num_inference_steps = int(inference_steps)
     texture_model.config.reference_conditioning_scale = float(reference_conditioning_scale)
     
-    # La semilla ya viene procesada
-    final_seed = int(seed)
-    print(f"Parámetros: Pasos={inference_steps}, Escala Guía={guidance_scale}, Semilla Final={final_seed}")
+    print(f"[{session_hash}] Parámetros: Pasos={inference_steps}, Escala Guía={guidance_scale}, Semilla={final_seed}")
 
-    print(f"Cargando malla desde: {input_mesh_path}")
+    print(f"[{session_hash}] Cargando malla desde: {input_mesh_path}")
     try:
         user_mesh = trimesh.load(input_mesh_path, force='mesh')
     except Exception as e:
         raise gr.Error(f"No se pudo cargar el archivo del modelo 3D. Error: {e}")
 
-    print("Pre-procesando la malla...")
+    print(f"[{session_hash}] Pre-procesando la malla...")
     user_mesh = remove_degenerate_face(user_mesh)
     user_mesh = reduce_face(user_mesh)
     
-    print("Ejecutando el pipeline de texturizado de Step1X-3D...")
+    print(f"[{session_hash}] Ejecutando el pipeline de texturizado...")
     textured_mesh = texture_model(
         image=input_image_path, 
         mesh=user_mesh, 
@@ -138,16 +185,16 @@ def generate_texture(
     )
     
     save_name = str(uuid.uuid4())
-    textured_save_path = f"{args.cache_dir}/{save_name}-textured.glb"
+    textured_save_path = f"{user_dir}/{save_name}-textured.glb"
     textured_mesh.export(textured_save_path)
     
     torch.cuda.empty_cache()
-    print(f"Malla texturizada guardada en: {textured_save_path}")
+    print(f"[{session_hash}] Malla texturizada guardada en: {textured_save_path}")
     
     return textured_save_path
 
 # --------------------------------------------------------------------------
-# 4. INTERFAZ DE USUARIO CON GRADIO
+# 6. INTERFAZ DE USUARIO CON GRADIO
 # --------------------------------------------------------------------------
 
 with gr.Blocks(title="Step1X-3D Texture Generator") as demo:
@@ -164,8 +211,8 @@ with gr.Blocks(title="Step1X-3D Texture Generator") as demo:
                 inference_steps = gr.Slider(label="Pasos de Inferencia (Calidad vs. Velocidad)", minimum=10, maximum=100, value=50, step=1)
                 reference_conditioning_scale = gr.Slider(label="Escala de Condicionamiento de Imagen", minimum=0.0, maximum=2.0, step=0.1, value=1.0)
                 with gr.Row():
-                    seed = gr.Slider(label="Semilla (Seed)", minimum=0, maximum=MAX_SEED, step=1, value=2024, interactive=True)
-                    randomize_seed = gr.Checkbox(label="Aleatoria", value=True)
+                    seed_input = gr.Slider(label="Semilla (Seed)", minimum=0, maximum=MAX_SEED, step=1, value=2024, interactive=True)
+                    randomize_seed_checkbox = gr.Checkbox(label="Aleatoria", value=True)
 
             btn_generate = gr.Button("✨ Generar Textura", variant="primary")
 
@@ -173,22 +220,25 @@ with gr.Blocks(title="Step1X-3D Texture Generator") as demo:
             output_model = gr.Model3D(label="Resultado: Modelo Texturizado", height=600, clear_color=[0.0, 0.0, 0.0, 0.0])
 
     # --- Lógica de la interfaz ---
+    
+    # Conecta las funciones de sesión al ciclo de vida de la demo
+    demo.load(start_session)
+    demo.unload(end_session)
 
+    # Conecta el botón a la función de generación
     btn_generate.click(
-        fn=get_random_seed,
-        inputs=[randomize_seed, seed],
-        outputs=[seed]
-    ).then(
-        fn=generate_texture,
+        fn=generate_texture_for_user_mesh,
         inputs=[
             input_image,
             input_mesh,
             guidance_scale,
             inference_steps,
-            seed, 
+            seed_input,
+            randomize_seed_checkbox,
             reference_conditioning_scale,
         ],
         outputs=[output_model]
     )
 
+# Lanza la aplicación
 demo.launch(ssr_mode=False)