Update app.py

app.py

@@ -2,73 +2,160 @@ 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
 
-# --- Instalación del Entorno (sin cambios) ---
+# --------------------------------------------------------------------------
+# 1. INSTALACIÓN DEL ENTORNO Y COMPILACIÓN DE EXTENSIONES
+# --------------------------------------------------------------------------
 def install_cuda_toolkit():
+    print("Verificando e instalando CUDA Toolkit si es necesario...")
     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)
-    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"])
+    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)
+    else:
+        print("CUDA Toolkit ya está instalado.")
 
     os.environ["CUDA_HOME"] = "/usr/local/cuda"
-    os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
-    os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
-        os.environ["CUDA_HOME"],
-        "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
-    )
+    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"
 
 install_cuda_toolkit()
+print("--- Versiones de librerías ---")
 os.system("pip list | grep torch")
 os.system('nvcc -V')
-print("cd /home/user/app/step1x3d_texture/differentiable_renderer/ && python setup.py install")
-os.system("cd /home/user/app/step1x3d_texture/differentiable_renderer/ && python setup.py install")
+print("-----------------------------")
 
-subprocess.run(shlex.split("pip install custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl"), check=True)
+# --- 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")
 
-import time
-import uuid
-import torch
-import trimesh
-import argparse
-import numpy as np
-import gradio as gr
-from step1x3d_geometry.models.pipelines.pipeline import Step1X3DGeometryPipeline
-from step1x3d_texture.pipelines.step1x_3d_texture_synthesis_pipeline import (
-    Step1X3DTexturePipeline,
+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.")
+
+
+# --------------------------------------------------------------------------
+# 2. IMPORTACIONES PRINCIPALES Y CARGA DE MODELOS
+# --------------------------------------------------------------------------
+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
 
-# --- Configuración y Carga de Modelos (sin cambios) ---
 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("--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"
+print(f"Usando dispositivo: {device}")
 
+# --- 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
 geometry_model = Step1X3DGeometryPipeline.from_pretrained(
     "stepfun-ai/Step1X-3D", subfolder=args.geometry_model
-).to("cuda")
+).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
 texture_model = Step1X3DTexturePipeline.from_pretrained("stepfun-ai/Step1X-3D", subfolder=args.texture_model)
+print("Pipeline de texturizado cargado.")
 
 
-# --- NUEVAS FUNCIONES DE GENERACIÓN SEPARADAS ---
+# --- 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)
+pipe_control = StableDiffusionXLControlNetPipeline.from_pretrained(
+    "sd-community/sdxl-flash", controlnet=controlnet, vae=vae, torch_dtype=torch.float16
+)
+pipe_control.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_control.scheduler.config)
+pipe_control.to(device)
+print("Modelos de ControlNet cargados y listos.")
+
+
+# --------------------------------------------------------------------------
+# 3. FUNCIONES DE BACKEND PARA CADA ETAPA
+# --------------------------------------------------------------------------
+
+@spaces.GPU(duration=60)
+def enhance_sketch(image, prompt):
+    """
+    Paso 0: Convierte un boceto en una imagen de alta calidad usando ControlNet.
+    """
+    if image is None:
+        raise gr.Error("Por favor, proporciona un boceto de entrada.")
+
+    input_image = image.convert("RGB")
+    
+    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"
+    
+    print(f"Mejorando boceto con prompt: '{final_prompt}'")
+    output_image = pipe_control(
+        prompt=final_prompt,
+        negative_prompt=negative_prompt,
+        image=input_image,
+        num_inference_steps=20,
+        controlnet_conditioning_scale=0.85,
+        guidance_scale=5.0,
+        width=new_width,
+        height=new_height,
+    ).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
 
 @spaces.GPU(duration=180)
-def generate_geometry(
-    input_image_path, guidance_scale, inference_steps, max_facenum, symmetry, edge_type
-):
+def generate_geometry(input_image_path, guidance_scale, inference_steps, max_facenum, symmetry, edge_type):
     """
-    Función que genera únicamente la geometría.
+    Paso 1: Genera la geometría a partir de una imagen.
     """
+    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.")
+        
     print("Iniciando generación de geometría...")
     if "Label" in args.geometry_model:
         symmetry_values = ["x", "asymmetry"]
@@ -95,26 +182,24 @@ 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
     return geometry_save_path, geometry_save_path
 
 @spaces.GPU(duration=120)
 def generate_texture(input_image_path, geometry_path):
     """
-    Función que aplica la textura a una geometría existente.
+    Paso 2: Aplica la textura a una geometría existente.
     """
     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.")
         
     print(f"Iniciando texturizado para la malla: {geometry_path}")
     geometry_mesh = trimesh.load(geometry_path)
     
-    # Post-procesamiento de la geometría antes de texturizar
     geometry_mesh = remove_degenerate_face(geometry_mesh)
     geometry_mesh = reduce_face(geometry_mesh)
     
-    # Aplicar textura
     textured_mesh = texture_model(input_image_path, geometry_mesh)
     
     save_name = os.path.basename(geometry_path).replace(".glb", "")
@@ -123,91 +208,75 @@ def generate_texture(input_image_path, geometry_path):
     
     torch.cuda.empty_cache()
     print(f"Malla texturizada guardada en: {textured_save_path}")
-    
     return textured_save_path
 
+# --------------------------------------------------------------------------
+# 4. INTERFAZ DE USUARIO CON GRADIO
+# --------------------------------------------------------------------------
 
-# --- NUEVA INTERFAZ DE GRADIO ---
-
-with gr.Blocks(title="Step1X-3D demo") as demo:
-    gr.Markdown("# Step1X-3D")
+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.")
     
-    # Estado para almacenar la ruta de la geometría generada
+    processed_image_path_state = gr.State()
     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",
-            )
+            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")
+                edge_type = gr.Radio(choices=["sharp", "normal", "smooth"], label="Edge Type", value="sharp", type="value")
             
             with gr.Row():
-                btn_geo = gr.Button("1. Generate Geometry")
-                btn_tex = gr.Button("2. Generate Texture", visible=False)
+                btn_enhance = gr.Button("Paso 0: Mejorar Boceto")
+            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):
-            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])
+            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])
             
         with gr.Column(scale=1):
             gr.Examples(
                 examples=[
-                    ["examples/images/000.png"],
-                    ["examples/images/001.png"],
-                    ["examples/images/004.png"],
-                    ["examples/images/008.png"],
-                    ["examples/images/028.png"],
-                    ["examples/images/032.png"],
-                    ["examples/images/061.png"],
-                    ["examples/images/107.png"],
+                    [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"],
                 ],
-                inputs=[input_image],
-                cache_examples=False,
+                inputs=[input_image, prompt],
+                cache_examples=False
             )
 
-    # Lógica de los botones
+    # --- 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])
     
-    # Al hacer clic en el botón de geometría
-    btn_geo.click(
-        fn=generate_geometry,
-        inputs=[
-            input_image,
-            guidance_scale,
-            inference_steps,
-            max_facenum,
-            symmetry,
-            edge_type,
-        ],
-        outputs=[geometry_preview, geometry_path_state]
-    ).then(
-        # Cuando la geometría termine, actualiza la interfaz
-        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
-        },
-        outputs=[btn_tex, textured_preview]
+    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]
     )
     
-    # Al hacer clic en el botón de textura
-    btn_tex.click(
-        fn=generate_texture,
-        inputs=[input_image, geometry_path_state],
-        outputs=[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_tex.click(fn=generate_texture, inputs=[processed_image_path_state, geometry_path_state], outputs=[textured_preview])
 
-demo.launch(ssr_mode=False)
+demo.launch(ssr_mode=False, share=True)