Update app.py

app.py

@@ -1,75 +1,62 @@
+# ==============================================================================
+# 1. INSTALACIÓN DEL ENTORNO Y DEPENDENCIAS
+# ==============================================================================
 import os
 import shlex
 import spaces
 import subprocess
-import time
-import uuid
-import torch
-import trimesh
-import argparse
-import numpy as np
-import gradio as gr
-from PIL import Image, ImageOps
+import logging
+
+# Configuración del logging para depuración
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - Step1X-3D - %(levelname)s - %(message)s')
 
-# --------------------------------------------------------------------------
-# 1. INSTALACIÓN DEL ENTORNO Y COMPILACIÓN DE EXTENSIONES
-# --------------------------------------------------------------------------
-def install_cuda_toolkit():
-    print("Verificando e instalando CUDA Toolkit si es necesario...")
+def install_dependencies():
+    """Instala el toolkit de CUDA y compila las extensiones C++/CUDA necesarias."""
+    logging.info("Iniciando la instalación de dependencias...")
+    
+    # Instalar 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)}"
-    if not os.path.exists("/usr/local/cuda/bin/nvcc"):
-        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"], check=True)
+    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"])
     else:
-        print("CUDA Toolkit ya está instalado.")
+        logging.info("CUDA Toolkit ya está instalado.")
 
     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', '')}"
     os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
+    
+    # Compilar extensiones personalizadas
+    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.")
+    os.system('nvcc -V')
 
-install_cuda_toolkit()
-print("--- Versiones de librerías ---")
-os.system("pip list | grep torch")
-os.system('nvcc -V')
-print("-----------------------------")
-
-# --- RUTAS CORREGIDAS ---
-# Construir rutas absolutas basadas en la ubicación del script app.py
-APP_DIR = os.path.dirname(os.path.abspath(__file__))
-PROJECT_ROOT = os.path.join(APP_DIR, "Step1X-3D")
-DIFFERENTIABLE_RENDERER_PATH = os.path.join(PROJECT_ROOT, "step1x3d_texture/differentiable_renderer/")
-CUSTOM_RASTERIZER_WHL_PATH = os.path.join(PROJECT_ROOT, "custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl")
-
-print(f"Ruta del renderizador diferenciable: {DIFFERENTIABLE_RENDERER_PATH}")
-print(f"Ruta del rasterizador personalizado (.whl): {CUSTOM_RASTERIZER_WHL_PATH}")
-
-
-print("Compilando extensión C++/CUDA para el renderizador diferenciable...")
-# Usamos cwd (current working directory) para ejecutar el comando en la carpeta correcta
-subprocess.run(
-    "python setup.py install",
-    shell=True, check=True, cwd=DIFFERENTIABLE_RENDERER_PATH
-)
-
-print("Instalando la extensión C++/CUDA para el rasterizador personalizado...")
-subprocess.run(
-    shlex.split(f"pip install {CUSTOM_RASTERIZER_WHL_PATH}"),
-    check=True
-)
-print("Compilaciones finalizadas.")
-
+install_dependencies()
 
-# --------------------------------------------------------------------------
-# 2. IMPORTACIONES PRINCIPALES Y CARGA DE MODELOS
-# --------------------------------------------------------------------------
+import uuid
+import torch
+import trimesh
+import argparse
+import numpy as np
+import gradio as gr
+from PIL import Image, ImageOps
 from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL, EulerAncestralDiscreteScheduler
 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
+# ==============================================================================
+
 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")
@@ -78,40 +65,43 @@ args = parser.parse_args()
 
 os.makedirs(args.cache_dir, exist_ok=True)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Usando dispositivo: {device}")
+torch_dtype = torch.float16
+
+logging.info("Cargando modelos... Este proceso puede tardar varios minutos.")
 
-# --- Carga de Modelos Step1X-3D ---
-print("Cargando pipeline de geometría Step1X-3D...")
-geometry_model_path = os.path.join(PROJECT_ROOT) # La ruta base para from_pretrained
+# Carga de modelos de 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("Pipeline de geometría cargado.")
 
-print("Cargando pipeline de texturizado Step1X-3D...")
-texture_model_path = os.path.join(PROJECT_ROOT) # La ruta base para from_pretrained
+logging.info(f"Cargando modelo de textura: {args.texture_model}")
 texture_model = Step1X3DTexturePipeline.from_pretrained("stepfun-ai/Step1X-3D", subfolder=args.texture_model)
-print("Pipeline de texturizado cargado.")
-
 
-# --- Carga de Modelos de ControlNet para el pre-procesamiento ---
-print("Cargando modelos para el pre-procesamiento de bocetos (ControlNet)...")
-controlnet = ControlNetModel.from_pretrained("xinsir/controlnet-scribble-sdxl-1.0", torch_dtype=torch.float16)
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+# Carga de modelos de ControlNet para el pre-procesamiento de bocetos
+logging.info("Cargando modelos para el pre-procesamiento de bocetos (SDXL + ControlNet)...")
+controlnet = ControlNetModel.from_pretrained("xinsir/controlnet-scribble-sdxl-1.0", torch_dtype=torch_dtype)
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype)
 pipe_control = StableDiffusionXLControlNetPipeline.from_pretrained(
-    "sd-community/sdxl-flash", controlnet=controlnet, vae=vae, torch_dtype=torch.float16
+    "sd-community/sdxl-flash", controlnet=controlnet, vae=vae, torch_dtype=torch_dtype
 )
 pipe_control.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_control.scheduler.config)
 pipe_control.to(device)
-print("Modelos de ControlNet cargados y listos.")
+logging.info("Todos los modelos han sido cargados correctamente.")
 
 
-# --------------------------------------------------------------------------
-# 3. FUNCIONES DE BACKEND PARA CADA ETAPA
-# --------------------------------------------------------------------------
+# ==============================================================================
+# 3. FUNCIONES DE GENERACIÓN POR PASOS
+# ==============================================================================
+
+def apply_3d_style(prompt: str) -> tuple[str, str]:
+    """Aplica el estilo '3D Model' por defecto al prompt."""
+    style_prompt = "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting"
+    negative_prompt = "ugly, deformed, noisy, low poly, blurry, painting"
+    return style_prompt.replace("{prompt}", prompt), negative_prompt
 
 @spaces.GPU(duration=60)
-def enhance_sketch(image, prompt):
+def process_sketch(image, prompt, negative_prompt, guidance_scale, num_steps, controlnet_scale):
     """
     Paso 0: Convierte un boceto en una imagen de alta calidad usando ControlNet.
     """
@@ -120,43 +110,48 @@ def enhance_sketch(image, prompt):
 
     input_image = image.convert("RGB")
     
+    # Pre-procesamiento de la imagen de entrada (invertir y redimensionar)
     width, height = input_image.size
     ratio = np.sqrt(1024.0 * 1024.0 / (width * height))
     new_width, new_height = int(width * ratio), int(height * ratio)
     input_image = input_image.resize((new_width, new_height))
     input_image = ImageOps.invert(input_image)
     
-    final_prompt = f"professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting"
-    negative_prompt = "ugly, deformed, noisy, low poly, blurry, painting"
+    final_prompt, final_negative_prompt = apply_3d_style(prompt)
+    if negative_prompt: # Añadir negativo del usuario si existe
+        final_negative_prompt = f"{final_negative_prompt}, {negative_prompt}"
+
+    logging.info(f"Mejorando boceto con prompt: '{final_prompt}'")
     
-    print(f"Mejorando boceto con prompt: '{final_prompt}'")
+    generator = torch.Generator(device=device).manual_seed(np.random.randint(0, 2**32 - 1))
+
     output_image = pipe_control(
         prompt=final_prompt,
-        negative_prompt=negative_prompt,
+        negative_prompt=final_negative_prompt,
         image=input_image,
-        num_inference_steps=20,
-        controlnet_conditioning_scale=0.85,
-        guidance_scale=5.0,
+        num_inference_steps=int(num_steps),
+        controlnet_conditioning_scale=float(controlnet_scale),
+        guidance_scale=float(guidance_scale),
         width=new_width,
         height=new_height,
+        generator=generator,
     ).images[0]
     
     save_name = str(uuid.uuid4())
     processed_image_path = f"{args.cache_dir}/{save_name}_processed.png"
     output_image.save(processed_image_path)
     
-    print(f"Boceto mejorado y guardado en: {processed_image_path}")
-    return processed_image_path, processed_image_path
+    logging.info(f"Boceto mejorado y guardado en: {processed_image_path}")
+    return processed_image_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 una imagen.
-    """
+    """Paso 1: Genera la geometría a partir de la imagen procesada."""
     if not input_image_path or not os.path.exists(input_image_path):
-        raise gr.Error("Primero debes proporcionar una imagen o mejorar un boceto.")
+        raise gr.Error("Primero debes procesar un boceto o proporcionar una imagen de entrada válida.")
         
-    print("Iniciando generación de geometría...")
+    logging.info(f"Iniciando generación de geometría desde: {os.path.basename(input_image_path)}")
+    
     if "Label" in args.geometry_model:
         symmetry_values = ["x", "asymmetry"]
         out = geometry_model(
@@ -175,108 +170,143 @@ def generate_geometry(input_image_path, guidance_scale, inference_steps, max_fac
             max_facenum=int(max_facenum),
         )
 
-    save_name = str(uuid.uuid4())
-    geometry_save_path = f"{args.cache_dir}/{save_name}.glb"
+    save_name = os.path.basename(input_image_path).replace("_processed.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()
-    print(f"Geometría guardada en: {geometry_save_path}")
-    return geometry_save_path, geometry_save_path
+    logging.info(f"Geometría guardada en: {geometry_save_path}")
+    return geometry_save_path
 
 @spaces.GPU(duration=120)
 def generate_texture(input_image_path, geometry_path):
-    """
-    Paso 2: Aplica la textura a una geometría existente.
-    """
+    """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 antes de texturizar.")
     if not input_image_path or not os.path.exists(input_image_path):
-        raise gr.Error("Se necesita una imagen de entrada para el texturizado.")
+        raise gr.Error("Se necesita la imagen procesada para el texturizado.")
         
-    print(f"Iniciando texturizado para la malla: {geometry_path}")
+    logging.info(f"Iniciando texturizado para la malla: {os.path.basename(geometry_path)}")
     geometry_mesh = trimesh.load(geometry_path)
     
+    # Post-procesamiento
     geometry_mesh = remove_degenerate_face(geometry_mesh)
     geometry_mesh = reduce_face(geometry_mesh)
     
     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"
+    save_name = os.path.basename(geometry_path).replace("_geometry.glb", "")
+    textured_save_path = f"{args.cache_dir}/{save_name}_textured.glb"
     textured_mesh.export(textured_save_path)
     
     torch.cuda.empty_cache()
-    print(f"Malla texturizada guardada en: {textured_save_path}")
+    logging.info(f"Malla texturizada guardada en: {textured_save_path}")
     return textured_save_path
 
-# --------------------------------------------------------------------------
-# 4. INTERFAZ DE USUARIO CON GRADIO
-# --------------------------------------------------------------------------
+# ==============================================================================
+# 4. INTERFAZ DE GRADIO
+# ==============================================================================
 
-with gr.Blocks(title="Step1X-3D: Flujo de Boceto a 3D") as demo:
-    gr.Markdown("# Step1X-3D: Flujo de Boceto a 3D")
-    gr.Markdown("Flujo de trabajo inspirado en TRELLIS. **Paso 0 (Opcional):** Sube un boceto y un prompt para generar una imagen de referencia. **Paso 1:** Genera la geometría 3D a partir de la imagen. **Paso 2:** Genera la textura para la geometría.")
+with gr.Blocks(title="Step1X-3D", css="footer {display: none !important;} a {text-decoration: none !important;}") as demo:
+    gr.Markdown("# Step1X-3D: De Boceto a Malla 3D Texturizada")
+    gr.Markdown("Flujo de trabajo en 3 pasos: **0. Procesar Boceto → 1. Generar Geometría → 2. Generar Textura**")
     
+    # Estados para mantener las rutas de los archivos entre pasos
     processed_image_path_state = gr.State()
     geometry_path_state = gr.State()
 
     with gr.Row():
         with gr.Column(scale=2):
-            gr.Markdown("### **Entrada Principal**")
-            input_image = gr.Image(label="Sube tu boceto o imagen de referencia", type="pil", image_mode="RGBA")
-            prompt = gr.Textbox(label="Prompt (describe tu objeto)", value="a detailed medieval sword")
-
-            with gr.Accordion(label="Opciones de Generación 3D", open=True):
-                guidance_scale = gr.Number(label="Guidance Scale (3D)", value="7.5")
-                inference_steps = gr.Slider(label="Inference Steps (3D)", minimum=1, maximum=100, value=50, step=1)
-                max_facenum = gr.Number(label="Max Face Num", value="400000")
-                symmetry = gr.Radio(choices=["Symmetry", "Asymmetry"], label="Symmetry", value="Symmetry", type="index")
+            # --- Panel de Entradas ---
+            input_image = gr.Image(label="Paso 0: Carga tu boceto o imagen", type="pil", image_mode="RGB")
+            prompt = gr.Textbox(label="Describe tu objeto", value="a comfortable armchair")
+
+            with gr.Accordion(label="Opciones Avanzadas", open=False):
+                gr.Markdown("### Opciones de Procesado de Boceto (Paso 0)")
+                neg_prompt_sketch = gr.Textbox(label="Negative Prompt (Boceto)", value="text, signature, watermark")
+                guidance_sketch = gr.Slider(0.1, 10.0, label="Guidance Scale (Boceto)", value=5.0, step=0.1)
+                steps_sketch = gr.Slider(1, 50, label="Steps (Boceto)", value=25, step=1)
+                controlnet_scale = gr.Slider(0.1, 2.0, label="ControlNet Scale", value=0.85, step=0.05)
+
+                gr.Markdown("---")
+                gr.Markdown("### Opciones de Generación 3D (Paso 1)")
+                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_enhance = gr.Button("Paso 0: Mejorar Boceto")
+                btn_process_sketch = gr.Button("0. Procesar Boceto", variant="secondary")
             with gr.Row():
-                btn_geo = gr.Button("Paso 1: Generar Geometría", interactive=False)
-                btn_tex = gr.Button("Paso 2: Generar Textura", visible=False, interactive=False)
-
-        with gr.Column(scale=4):
-            gr.Markdown("### **Resultados**")
-            processed_image_preview = gr.Image(label="Boceto Mejorado (Entrada para 3D)", type="filepath", interactive=False)
-            textured_preview = gr.Model3D(label="Modelo Texturizado Final", height=380, clear_color=[0.0, 0.0, 0.0, 0.0])
-            geometry_preview = gr.Model3D(label="Modelo Geométrico", height=380, clear_color=[0.0, 0.0, 0.0, 0.0])
+                btn_geo = gr.Button("1. Generar Geometría", interactive=False)
+                btn_tex = gr.Button("2. Generar Textura", interactive=False)
+
+        with gr.Column(scale=3):
+            # --- Panel de Salidas ---
+            processed_image_preview = gr.Image(label="Resultado del Boceto Procesado", 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])
             
         with gr.Column(scale=1):
             gr.Examples(
                 examples=[
-                    [os.path.join(PROJECT_ROOT, "examples/images/000.png"), "a toy car"],
-                    [os.path.join(PROJECT_ROOT, "examples/images/001.png"), "a red fire hydrant"],
-                    [os.path.join(PROJECT_ROOT, "examples/images/004.png"), "a detailed, ornate, golden picture frame"],
+                    ["examples/images/000.png", "a futuristic spaceship"],
+                    ["examples/images/001.png", "a cartoon style monster"],
+                    ["examples/images/004.png", "a red sports car"],
+                    ["examples/images/008.png", "a medieval sword"],
+                    ["examples/images/028.png", "a vintage camera"],
+                    ["examples/images/032.png", "a cute robot"],
+                    ["examples/images/061.png", "a delicious hamburger"],
+                    ["examples/images/107.png", "a golden trophy"],
                 ],
-                inputs=[input_image, prompt],
-                cache_examples=False
+                inputs=[input_image, prompt], cache_examples=False
             )
 
     # --- Lógica de la Interfaz ---
 
-    def handle_image_upload(img):
-        if img is None:
-            return None, gr.update(interactive=False)
-        save_name = str(uuid.uuid4())
-        original_path = f"{args.cache_dir}/{save_name}_original.png"
-        img.save(original_path)
-        return original_path, gr.update(interactive=True)
-
-    input_image.upload(fn=handle_image_upload, inputs=[input_image], outputs=[processed_image_path_state, btn_geo])
-    
-    btn_enhance.click(fn=enhance_sketch, inputs=[input_image, prompt], outputs=[processed_image_preview, processed_image_path_state]).then(
-        fn=lambda: gr.update(interactive=True), outputs=[btn_geo]
+    def on_sketch_processed(path):
+        """Función a ejecutar cuando el boceto se ha procesado."""
+        return {
+            processed_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):
+        """Función a ejecutar cuando la geometría se ha generado."""
+        return {
+            geometry_path_state: path,
+            btn_tex: gr.update(interactive=True, variant="primary"),
+        }
+
+    btn_process_sketch.click(
+        fn=process_sketch,
+        inputs=[input_image, prompt, neg_prompt_sketch, guidance_sketch, steps_sketch, controlnet_scale],
+        outputs=[processed_image_preview]
+    ).then(
+        fn=on_sketch_processed,
+        inputs=[processed_image_preview],
+        outputs=[processed_image_path_state, btn_geo, btn_tex, geometry_preview, textured_preview]
     )
     
-    btn_geo.click(fn=generate_geometry, inputs=[processed_image_path_state, guidance_scale, inference_steps, max_facenum, symmetry, edge_type], outputs=[geometry_preview, geometry_path_state]).then(
-        fn=lambda: {btn_tex: gr.update(visible=True, interactive=True), textured_preview: gr.update(value=None)}, outputs=[btn_tex, textured_preview]
+    btn_geo.click(
+        fn=generate_geometry,
+        inputs=[processed_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=[processed_image_path_state, geometry_path_state], outputs=[textured_preview])
+    btn_tex.click(
+        fn=generate_texture,
+        inputs=[processed_image_path_state, geometry_path_state],
+        outputs=[textured_preview],
+    )
 
-demo.launch(ssr_mode=False, share=True)
+demo.launch(ssr_mode=False)