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-T2I - %(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...")
    
    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"
    
    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 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
# ==============================================================================

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)

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

    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"

    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]
    
    save_name = str(uuid.uuid4())
    image_save_path = f"{args.cache_dir}/{save_name}_t2i.png"
    image.save(image_save_path)
    
    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)}")
    
    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("_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

@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)}")
    
    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)
    
    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 (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()

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

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

demo.launch(ssr_mode=False)