app.py

# ==============================================================================
# 1. INSTALACIÓN DEL ENTORNO Y DEPENDENCIAS
# ==============================================================================
import os
import shlex
import spaces
import subprocess
import logging

# Configuración del logging para depuración
logging.basicConfig(level=logging.INFO, format='%(asctime)s - Step1X-3D - %(levelname)s - %(message)s')

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"):
        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:
        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_dependencies()

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

logging.info("Cargando modelos... Este proceso puede tardar varios minutos.")

# 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)

logging.info(f"Cargando modelo de textura: {args.texture_model}")
texture_model = Step1X3DTexturePipeline.from_pretrained("stepfun-ai/Step1X-3D", subfolder=args.texture_model)

# 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_dtype
)
pipe_control.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_control.scheduler.config)
pipe_control.to(device)
logging.info("Todos los modelos han sido cargados correctamente.")


# ==============================================================================
# 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 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.
    """
    if image is None:
        raise gr.Error("Por favor, proporciona un boceto de entrada.")

    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, 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}'")
    
    generator = torch.Generator(device=device).manual_seed(np.random.randint(0, 2**32 - 1))

    output_image = pipe_control(
        prompt=final_prompt,
        negative_prompt=final_negative_prompt,
        image=input_image,
        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)
    
    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 la imagen procesada."""
    if not input_image_path or not os.path.exists(input_image_path):
        raise gr.Error("Primero debes procesar un boceto o proporcionar una imagen de entrada válida.")
        
    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(
            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),
        )

    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()
    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 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 la imagen procesada para el texturizado.")
        
    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("_geometry.glb", "")
    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}")
    return textured_save_path

# ==============================================================================
# 4. INTERFAZ DE GRADIO
# ==============================================================================

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):
            # --- 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_process_sketch = gr.Button("0. Procesar Boceto", variant="secondary")
            with gr.Row():
                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=[
                    ["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
            )

    # --- Lógica de la Interfaz ---

    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_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],
    )

demo.launch(ssr_mode=False)