app.py

import os
import shlex
import spaces
import subprocess
import argparse
import uuid
import random
import numpy as np
import shutil

# --------------------------------------------------------------------------
# BLOQUE 1: INSTALACIÓN DEL ENTORNO Y DEPENDENCIAS (Sin cambios)
# --------------------------------------------------------------------------

@spaces.GPU(duration=120)
def setup_environment():
    """
    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, capture_output=True, text=True
        )
        print("Renderizador diferenciable compilado.")
    except subprocess.CalledProcessError as 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 ---")

setup_environment()

# --------------------------------------------------------------------------
# BLOQUE 2: IMPORTACIONES DE LA APLICACIÓN Y LÓGICA PRINCIPAL
# --------------------------------------------------------------------------
import torch
import trimesh
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()

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.")

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 y listo.")

# --------------------------------------------------------------------------
# 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):
    """Función auxiliar para determinar la semilla a usar."""
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    print(f"Usando semilla: {seed}")
    return seed

@spaces.GPU(duration=180)
def generate_texture(
    input_image_path,
    input_mesh_path,
    guidance_scale,
    inference_steps,
    seed,
    reference_conditioning_scale,
    req: gr.Request,
    progress=gr.Progress(track_tqdm=True)
):
    """
    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.")

    session_hash = str(req.session_hash)
    user_dir = os.path.join(TMP_DIR, session_hash)
    print(f"[{session_hash}] Iniciando generación de textura...")
    
    final_seed = int(seed)
    
    # 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)
    
    print(f"[{session_hash}] Parámetros: Pasos={inference_steps}, Escala Guía={guidance_scale}, Semilla={final_seed}")

    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(f"[{session_hash}] Pre-procesando la malla...")
    user_mesh = remove_degenerate_face(user_mesh)
    user_mesh = reduce_face(user_mesh)
    
    print(f"[{session_hash}] Ejecutando el pipeline de texturizado...")
    textured_mesh = texture_model(
        image=input_image_path, 
        mesh=user_mesh, 
        remove_bg=True,
        seed=final_seed
    )
    
    save_name = str(uuid.uuid4())
    textured_save_path = f"{user_dir}/{save_name}-textured.glb"
    textured_mesh.export(textured_save_path)
    
    torch.cuda.empty_cache()
    print(f"[{session_hash}] Malla texturizada guardada en: {textured_save_path}")
    
    return textured_save_path


# --------------------------------------------------------------------------
# 6. INTERFAZ DE USUARIO CON GRADIO
# --------------------------------------------------------------------------

with gr.Blocks(title="Step1X-3D Texture Generator") as demo:
    gr.Markdown("# Step1X-3D: Generador de Texturas para Mallas Personalizadas")
    gr.Markdown("### Sube tu propia imagen y modelo 3D (`.glb`, `.obj`) para generar una textura de alta calidad.")
    
    with gr.Row():
        with gr.Column(scale=2):
            input_image = gr.Image(label="Imagen de Referencia", type="filepath", sources=["upload"])
            input_mesh = gr.Model3D(label="Modelo 3D de Entrada (.glb, .obj)")
            
            with gr.Accordion(label="Parámetros Avanzados", open=True):
                guidance_scale = gr.Slider(label="Guidance Scale (Fidelidad a la imagen)", minimum=1.0, maximum=10.0, value=3.0, step=0.1)
                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_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")

        with gr.Column(scale=3):
            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 ---
    
    demo.load(fn=start_session)
    demo.unload(fn=end_session)

    btn_generate.click(
        fn=get_random_seed,
        inputs=[randomize_seed_checkbox, seed_input],
        outputs=[seed_input] 
    ).then(
        fn=generate_texture,
        inputs=[
            input_image,
            input_mesh,
            guidance_scale,
            inference_steps,
            seed_input, 
            reference_conditioning_scale,
        ],
        outputs=[output_model]
    )

# Lanza la aplicación
demo.launch(ssr_mode=False)