Update app.py

app.py

@@ -1,280 +1,214 @@
-# ==============================================================================
-# 1. INSTALACIÓN DEL ENTORNO Y DEPENDENCIAS
-# ==============================================================================
 import os
 import shlex
 import spaces
 import subprocess
-import logging
+import argparse
+import uuid
+import random
+import numpy as np
 
-# Configuración del logging para depuración
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - Step1X-3D-T2I - %(levelname)s - %(message)s')
+# --------------------------------------------------------------------------
+# BLOQUE 1: INSTALACIÓN DEL ENTORNO Y DEPENDENCIAS
+# Esta sección DEBE ejecutarse en su totalidad ANTES de importar torch o gradio.
+# --------------------------------------------------------------------------
 
-def install_dependencies():
-    """Instala el toolkit de CUDA y compila las extensiones C++/CUDA necesarias."""
-    logging.info("Iniciando la instalación de dependencias...")
-    
-    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)}"
+@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"):
-        logging.info("Descargando e instalando CUDA Toolkit...")
-        subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
-        subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
-        subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
+        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)
+        # Se añade --override para evitar problemas en algunos entornos de HF Spaces
+        subprocess.run([CUDA_TOOLKIT_FILE, "--silent", "--toolkit", "--override"], check=True)
+        print("CUDA Toolkit instalado.")
     else:
-        logging.info("CUDA Toolkit ya está instalado.")
+        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']}/lib:{os.environ.get('LD_LIBRARY_PATH', '')}"
+    # La ruta correcta en la mayoría de las distribuciones de Linux es lib64
+    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"
-    
-    logging.info("Compilando extensiones de renderizado...")
-    renderer_path = "/home/user/app/step1x3d_texture/differentiable_renderer/"
-    subprocess.run(f"cd {renderer_path} && python setup.py install", shell=True, check=True)
-    subprocess.run(shlex.split("pip install custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl"), check=True)
-    
-    logging.info("Instalación completada.")
+    print("Variables de entorno de CUDA configuradas.")
+
+    # 3. Verificar NVCC
+    print("Verificando versión de NVCC:")
     os.system('nvcc -V')
 
-install_dependencies()
+    # 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.")
 
-import uuid
+    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()
+
+
+# --------------------------------------------------------------------------
+# BLOQUE 2: IMPORTACIONES DE LA APLICACIÓN Y LÓGICA PRINCIPAL
+# Todas las importaciones que dependen de CUDA se hacen AHORA.
+# --------------------------------------------------------------------------
 import torch
 import trimesh
-import argparse
-import random
-import numpy as np
 import gradio as gr
-from PIL import Image
-from huggingface_hub import hf_hub_download
-from diffusers import FluxPipeline, FluxTransformer2DModel, GGUFQuantizationConfig
-from transformers import T5EncoderModel, BitsAndBytesConfig as BitsAndBytesConfigTF
-from step1x3d_geometry.models.pipelines.pipeline import Step1X3DGeometryPipeline
 from step1x3d_texture.pipelines.step1x_3d_texture_synthesis_pipeline import Step1X3DTexturePipeline
 from step1x3d_geometry.models.pipelines.pipeline_utils import reduce_face, remove_degenerate_face
 
-# ==============================================================================
-# 2. CONFIGURACIÓN Y CARGA DE MODELOS
-# ==============================================================================
-
+# --------------------------------------------------------------------------
+# 3. CONFIGURACIÓN Y CARGA DEL MODELO DE TEXTURA
+# --------------------------------------------------------------------------
 parser = argparse.ArgumentParser()
-parser.add_argument("--geometry_model", type=str, default="Step1X-3D-Geometry-Label-1300m")
 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)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.bfloat16
 MAX_SEED = np.iinfo(np.int32).max
 
-logging.info("Cargando modelos... Este proceso puede tardar varios minutos.")
-
-# --- Carga de Modelos Step1X-3D ---
-logging.info(f"Cargando modelo de geometría: {args.geometry_model}")
-geometry_model = Step1X3DGeometryPipeline.from_pretrained(
-    "stepfun-ai/Step1X-3D", subfolder=args.geometry_model
-).to(device)
+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.")
 
-logging.info(f"Cargando modelo de textura: {args.texture_model}")
+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.")
 
-# --- Carga de Modelo FLUX para Texto-a-Imagen ---
-logging.info("Cargando modelo FLUX para Texto-a-Imagen...")
-single_file_base_model = "camenduru/FLUX.1-dev-diffusers"
 
-# --- CORRECCIÓN AQUÍ ---
-# Descargar el archivo GGUF explícitamente usando hf_hub_download
-flux_repo_id = "gokaygokay/flux-game"
-flux_filename = "hyperflux_00001_.q8_0.gguf"
-logging.info(f"Descargando {flux_filename} desde {flux_repo_id}...")
-downloaded_flux_path = hf_hub_download(repo_id=flux_repo_id, filename=flux_filename)
-logging.info(f"Archivo GGUF descargado en: {downloaded_flux_path}")
-# --- FIN DE LA CORRECCIÓN ---
-
-quantization_config_tf = BitsAndBytesConfigTF(load_in_8bit=True, bnb_8bit_compute_dtype=torch_dtype)
-text_encoder_2 = T5EncoderModel.from_pretrained(single_file_base_model, subfolder="text_encoder_2", torch_dtype=torch_dtype, quantization_config=quantization_config_tf)
-
-# Usar la ruta local descargada
-transformer = FluxTransformer2DModel.from_single_file(downloaded_flux_path, subfolder="transformer", quantization_config=GGUFQuantizationConfig(compute_dtype=torch_dtype), torch_dtype=torch_dtype, config=single_file_base_model)
-flux_pipeline = FluxPipeline.from_pretrained(single_file_base_model, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch_dtype)
-flux_pipeline.to(device)
-logging.info("Todos los modelos han sido cargados correctamente.")
-
-
-# ==============================================================================
-# 3. FUNCIONES DE GENERACIÓN POR PASOS (Sin cambios)
-# ==============================================================================
-
-@spaces.GPU(duration=60)
-def generate_image_from_text(prompt, seed, randomize_seed, guidance_scale, num_steps):
-    """Paso 0: Genera una imagen 2D a partir de un prompt de texto usando FLUX."""
-    if not prompt:
-        raise gr.Error("El prompt de texto no puede estar vacío.")
+# --------------------------------------------------------------------------
+# 4. FUNCIÓN DE GENERACIÓN DE TEXTURA
+# --------------------------------------------------------------------------
 
+def get_random_seed(randomize_seed, seed):
+    """Genera una semilla aleatoria si se solicita."""
     if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    
-    logging.info(f"Generando imagen con prompt: '{prompt}' y seed: {seed}")
-    generator = torch.Generator(device=device).manual_seed(seed)
-    
-    full_prompt = f"professional 3d model {prompt}. octane render, highly detailed, volumetric, dramatic lighting, on a white background"
-    negative_prompt = "ugly, deformed, noisy, low poly, blurry, painting, photo, text, watermark"
+        return random.randint(0, MAX_SEED)
+    return int(seed)
 
-    image = flux_pipeline(
-        prompt=full_prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=float(guidance_scale),
-        num_inference_steps=int(num_steps),
-        width=1024,
-        height=1024,
-        generator=generator,
-    ).images[0]
+@spaces.GPU(duration=180)
+def generate_texture_for_user_mesh(
+    input_image_path,
+    input_mesh_path,
+    guidance_scale,
+    inference_steps,
+    seed,
+    randomize_seed,
+    reference_conditioning_scale,
+    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.")
+
+    print("Iniciando generación de textura para el modelo del usuario...")
     
-    save_name = str(uuid.uuid4())
-    image_save_path = f"{args.cache_dir}/{save_name}_t2i.png"
-    image.save(image_save_path)
+    seed = get_random_seed(randomize_seed, seed)
     
-    logging.info(f"Imagen 2D generada y guardada en: {image_save_path}")
-    return image_save_path
-
-@spaces.GPU(duration=180)
-def generate_geometry(input_image_path, guidance_scale, inference_steps, max_facenum, symmetry, edge_type):
-    """Paso 1: Genera la geometría a partir de la imagen generada."""
-    if not input_image_path or not os.path.exists(input_image_path):
-        raise gr.Error("Primero debes generar una imagen desde el texto.")
-        
-    logging.info(f"Iniciando generación de geometría desde: {os.path.basename(input_image_path)}")
+    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)
     
-    if "Label" in args.geometry_model:
-        symmetry_values = ["x", "asymmetry"]
-        out = geometry_model(
-            input_image_path,
-            label={"symmetry": symmetry_values[int(symmetry)], "edge_type": edge_type},
-            guidance_scale=float(guidance_scale),
-            octree_resolution=384,
-            max_facenum=int(max_facenum),
-            num_inference_steps=int(inference_steps),
-        )
-    else:
-        out = geometry_model(
-            input_image_path,
-            guidance_scale=float(guidance_scale),
-            num_inference_steps=int(inference_steps),
-            max_facenum=int(max_facenum),
-        )
+    print(f"Parámetros: Pasos={inference_steps}, Escala Guía={guidance_scale}, Semilla={seed}")
 
-    save_name = os.path.basename(input_image_path).replace("_t2i.png", "")
-    geometry_save_path = f"{args.cache_dir}/{save_name}_geometry.glb"
-    geometry_mesh = out.mesh[0]
-    geometry_mesh.export(geometry_save_path)
-    
-    torch.cuda.empty_cache()
-    logging.info(f"Geometría guardada en: {geometry_save_path}")
-    return geometry_save_path
+    print(f"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}")
 
-@spaces.GPU(duration=120)
-def generate_texture(input_image_path, geometry_path):
-    """Paso 2: Aplica la textura a la geometría generada."""
-    if not geometry_path or not os.path.exists(geometry_path):
-        raise gr.Error("Por favor, primero genera la geometría.")
-    if not input_image_path or not os.path.exists(input_image_path):
-        raise gr.Error("Se necesita la imagen generada para el texturizado.")
-        
-    logging.info(f"Iniciando texturizado para la malla: {os.path.basename(geometry_path)}")
+    print("Pre-procesando la malla...")
+    user_mesh = remove_degenerate_face(user_mesh)
+    user_mesh = reduce_face(user_mesh)
     
-    geometry_mesh = trimesh.load(geometry_path)
-    geometry_mesh = remove_degenerate_face(geometry_mesh)
-    geometry_mesh = reduce_face(geometry_mesh)
-    textured_mesh = texture_model(input_image_path, geometry_mesh)
+    print("Ejecutando el pipeline de texturizado de Step1X-3D...")
+    textured_mesh = texture_model(
+        image=input_image_path, 
+        mesh=user_mesh, 
+        remove_bg=True,
+        seed=seed
+    )
     
-    save_name = os.path.basename(geometry_path).replace("_geometry.glb", "")
-    textured_save_path = f"{args.cache_dir}/{save_name}_textured.glb"
+    save_name = str(uuid.uuid4())
+    textured_save_path = f"{args.cache_dir}/{save_name}-textured.glb"
     textured_mesh.export(textured_save_path)
     
     torch.cuda.empty_cache()
-    logging.info(f"Malla texturizada guardada en: {textured_save_path}")
+    print(f"Malla texturizada guardada en: {textured_save_path}")
+    
     return textured_save_path
 
-# ==============================================================================
-# 4. INTERFAZ DE GRADIO (Sin cambios)
-# ==============================================================================
 
-with gr.Blocks(title="Step1X-3D", css="footer {display: none !important;} a {text-decoration: none !important;}") as demo:
-    gr.Markdown("# Step1X-3D: Flujo de Texto a Malla 3D Texturizada")
-    gr.Markdown("Flujo de trabajo en 3 pasos: **0. Generar Imagen → 1. Generar Geometría → 2. Generar Textura**")
-    
-    generated_image_path_state = gr.State()
-    geometry_path_state = gr.State()
+# --------------------------------------------------------------------------
+# 5. 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):
-            prompt = gr.Textbox(label="Paso 0: Describe tu objeto", placeholder="Ej: a treasure chest, a sci-fi helmet, a cute dog")
+            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("Opciones de Generación de Imagen (Paso 0)", open=True):
-                seed = gr.Slider(0, MAX_SEED, label="Seed", value=42, step=1)
-                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-                guidance_t2i = gr.Slider(0.0, 10.0, label="Guidance Scale (Imagen)", value=3.5, step=0.1)
-                steps_t2i = gr.Slider(1, 20, label="Steps (Imagen)", value=8, step=1)
+            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 = gr.Slider(label="Semilla (Seed)", minimum=0, maximum=MAX_SEED, step=1, value=2024, interactive=True)
+                    randomize_seed = gr.Checkbox(label="Aleatoria", value=True)
 
-            with gr.Accordion("Opciones de Generación 3D (Pasos 1 y 2)", open=False):
-                guidance_3d = gr.Number(label="Guidance Scale (3D)", value="7.5")
-                steps_3d = gr.Slider(label="Inference Steps (3D)", minimum=1, maximum=100, value=50)
-                max_facenum = gr.Number(label="Max Face Num", value="200000")
-                symmetry = gr.Radio(choices=["symmetry", "asymmetry"], label="Symmetry", value="symmetry", type="index")
-                edge_type = gr.Radio(choices=["sharp", "normal", "smooth"], label="Edge Type", value="sharp", type="value")
-            
-            with gr.Row():
-                btn_t2i = gr.Button("0. Generate Image", variant="secondary")
-            with gr.Row():
-                btn_geo = gr.Button("1. Generate Geometry", interactive=False)
-                btn_tex = gr.Button("2. Generate Texture", interactive=False)
+            btn_generate = gr.Button("✨ Generar Textura", variant="primary")
 
         with gr.Column(scale=3):
-            generated_image_preview = gr.Image(label="Imagen Generada", type="filepath", interactive=False, height=400)
-            geometry_preview = gr.Model3D(label="Vista Previa de la Geometría", height=400, clear_color=[0.0, 0.0, 0.0, 0.0])
-            textured_preview = gr.Model3D(label="Vista Previa del Modelo Texturizado", height=400, clear_color=[0.0, 0.0, 0.0, 0.0])
-            
-    def on_image_generated(path):
-        return {
-            generated_image_path_state: path,
-            btn_geo: gr.update(interactive=True, variant="primary"),
-            btn_tex: gr.update(interactive=False),
-            geometry_preview: gr.update(value=None),
-            textured_preview: gr.update(value=None),
-        }
-
-    def on_geometry_generated(path):
-        return {
-            geometry_path_state: path,
-            btn_tex: gr.update(interactive=True, variant="primary"),
-        }
-
-    btn_t2i.click(
-        fn=generate_image_from_text,
-        inputs=[prompt, seed, randomize_seed, guidance_t2i, steps_t2i],
-        outputs=[generated_image_preview]
-    ).then(
-        fn=on_image_generated,
-        inputs=[generated_image_preview],
-        outputs=[generated_image_path_state, btn_geo, btn_tex, geometry_preview, textured_preview]
-    )
-    
-    btn_geo.click(
-        fn=generate_geometry,
-        inputs=[generated_image_path_state, guidance_3d, steps_3d, max_facenum, symmetry, edge_type],
-        outputs=[geometry_preview]
-    ).then(
-        fn=on_geometry_generated,
-        inputs=[geometry_preview],
-        outputs=[geometry_path_state, btn_tex]
-    )
-    
-    btn_tex.click(
-        fn=generate_texture,
-        inputs=[generated_image_path_state, geometry_path_state],
-        outputs=[textured_preview],
+            output_model = gr.Model3D(label="Resultado: Modelo Texturizado", height=600, clear_color=[0.0, 0.0, 0.0, 0.0])
+
+    btn_generate.click(
+        fn=generate_texture_for_user_mesh,
+        inputs=[
+            input_image,
+            input_mesh,
+            guidance_scale,
+            inference_steps,
+            seed,
+            randomize_seed,
+            reference_conditioning_scale,
+        ],
+        outputs=[output_model]
     )
 
+# Lanza la aplicación (share=True no es necesario en HF Spaces)
 demo.launch(ssr_mode=False)