app.py

import os
import sys
import shutil
import logging
import traceback
from typing import *

import gradio as gr
import spaces
import torch
import numpy as np
import imageio
from easydict import EasyDict as edict

from trellis.pipelines import TrellisTextTo3DPipeline
from trellis.representations import Gaussian, MeshExtractResult
from trellis.utils import render_utils, postprocessing_utils

# Configuración de entorno
os.environ["TOKENIZERS_PARALLELISM"] = "true"
os.environ["SPCONV_ALGO"] = "native"

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - HF_SPACE - %(levelname)s - %(message)s"
)

MAX_SEED = np.iinfo(np.int32).max
TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tmp")
os.makedirs(TMP_DIR, exist_ok=True)


# -----------------------------
# Funciones de manejo de sesión
# -----------------------------
def start_session(req: gr.Request):
    session_hash = str(req.session_hash)
    user_dir = os.path.join(TMP_DIR, session_hash)
    logging.info(f"START SESSION: Creando directorio para la sesión {session_hash} en {user_dir}")
    os.makedirs(user_dir, exist_ok=True)


def end_session(req: gr.Request):
    session_hash = str(req.session_hash)
    user_dir = os.path.join(TMP_DIR, session_hash)
    logging.info(f"END SESSION: Intentando eliminar el directorio de la sesión {session_hash} en {user_dir}")

    if os.path.exists(user_dir):
        try:
            shutil.rmtree(user_dir)
            logging.info(f"Directorio de la sesión {session_hash} eliminado correctamente.")
        except Exception as e:
            logging.error(f"Error al eliminar el directorio de la sesión {session_hash}: {e}")
    else:
        logging.warning(
            f"El directorio de la sesión {session_hash} no fue encontrado. "
            "Es posible que ya haya sido limpiado."
        )


# -----------------------------
# Manejo de estado
# -----------------------------
def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
    return {
        "gaussian": {
            **gs.init_params,
            "_xyz": gs._xyz.cpu().numpy(),
            "_features_dc": gs._features_dc.cpu().numpy(),
            "_scaling": gs._scaling.cpu().numpy(),
            "_rotation": gs._rotation.cpu().numpy(),
            "_opacity": gs._opacity.cpu().numpy(),
        },
        "mesh": {
            "vertices": mesh.vertices.cpu().numpy(),
            "faces": mesh.faces.cpu().numpy(),
        },
    }


def unpack_state(state: dict) -> Tuple[Gaussian, edict]:
    gs = Gaussian(
        aabb=state["gaussian"]["aabb"],
        sh_degree=state["gaussian"]["sh_degree"],
        mininum_kernel_size=state["gaussian"]["mininum_kernel_size"],
        scaling_bias=state["gaussian"]["scaling_bias"],
        opacity_bias=state["gaussian"]["opacity_bias"],
        scaling_activation=state["gaussian"]["scaling_activation"],
    )
    gs._xyz = torch.tensor(state["gaussian"]["_xyz"], device="cuda")
    gs._features_dc = torch.tensor(state["gaussian"]["_features_dc"], device="cuda")
    gs._scaling = torch.tensor(state["gaussian"]["_scaling"], device="cuda")
    gs._rotation = torch.tensor(state["gaussian"]["_rotation"], device="cuda")
    gs._opacity = torch.tensor(state["gaussian"]["_opacity"], device="cuda")

    mesh = edict(
        vertices=torch.tensor(state["mesh"]["vertices"], device="cuda"),
        faces=torch.tensor(state["mesh"]["faces"], device="cuda"),
    )

    return gs, mesh


# -----------------------------
# Funciones utilitarias
# -----------------------------
def get_seed(randomize_seed: bool, seed: int) -> int:
    new_seed = np.random.randint(0, MAX_SEED) if randomize_seed else seed
    logging.info(f"Usando seed: {new_seed}")
    return new_seed


# -----------------------------
# Procesos principales
# -----------------------------
@spaces.GPU
def text_to_3d(
    prompt: str,
    seed: int,
    ss_guidance_strength: float,
    ss_sampling_steps: int,
    slat_guidance_strength: float,
    slat_sampling_steps: int,
    req: gr.Request,
) -> Tuple[dict, str]:
    session_hash = str(req.session_hash)
    logging.info(f"[{session_hash}] Iniciando text_to_3d con prompt: '{prompt[:50]}...'")
    user_dir = os.path.join(TMP_DIR, session_hash)

    outputs = pipeline.run(
        prompt,
        seed=seed,
        formats=["gaussian", "mesh"],
        sparse_structure_sampler_params={
            "steps": ss_sampling_steps,
            "cfg_strength": ss_guidance_strength,
        },
        slat_sampler_params={
            "steps": slat_sampling_steps,
            "cfg_strength": slat_guidance_strength,
        },
    )

    logging.info(f"[{session_hash}] Generación completada. Renderizando video...")
    video = render_utils.render_video(outputs["gaussian"][0], num_frames=120)["color"]
    video_geo = render_utils.render_video(outputs["mesh"][0], num_frames=120)["normal"]

    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
    video_path = os.path.join(user_dir, "sample.mp4")
    imageio.mimsave(video_path, video, fps=15)

    state = pack_state(outputs["gaussian"][0], outputs["mesh"][0])
    torch.cuda.empty_cache()

    logging.info(f"[{session_hash}] Video y estado listos. Devolviendo: {video_path}")
    return state, video_path


@spaces.GPU(duration=90)
def extract_glb(
    state: dict,
    mesh_simplify: float,
    texture_size: int,
    req: gr.Request,
) -> Tuple[str, str]:
    session_hash = str(req.session_hash)
    logging.info(f"[{session_hash}] Iniciando extract_glb...")
    user_dir = os.path.join(TMP_DIR, session_hash)

    gs, mesh = unpack_state(state)
    glb = postprocessing_utils.to_glb(
        gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False
    )

    glb_path = os.path.join(user_dir, "sample.glb")
    glb.export(glb_path)

    torch.cuda.empty_cache()
    logging.info(f"[{session_hash}] GLB listo: {glb_path}")
    return glb_path, glb_path


@spaces.GPU
def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    gs, _ = unpack_state(state)
    gaussian_path = os.path.join(user_dir, "sample.ply")
    gs.save_ply(gaussian_path)
    torch.cuda.empty_cache()
    return gaussian_path, gaussian_path


# -----------------------------
# Interfaz Gradio
# -----------------------------
with gr.Blocks(delete_cache=(600, 600)) as demo:
    gr.Markdown("""
    # UTPL - Conversión de Texto a objetos 3D usando IA  
    ### Tesis: *"Objetos tridimensionales creados por IA: Innovación en entornos virtuales"*  
    **Autor:** Carlos Vargas  
    **Base técnica:** Adaptación de [TRELLIS](https://trellis3d.github.io/)  
    **Propósito educativo:** Demostraciones académicas e investigación en modelado 3D automático  
    """)

    with gr.Row():
        with gr.Column():
            text_prompt = gr.Textbox(label="Text Prompt", lines=5)

            with gr.Accordion(label="Generation Settings", open=False):
                seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)

                gr.Markdown("Stage 1: Sparse Structure Generation")
                with gr.Row():
                    ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
                    ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=25, step=1)

                gr.Markdown("Stage 2: Structured Latent Generation")
                with gr.Row():
                    slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
                    slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=25, step=1)

            generate_btn = gr.Button("Generate")

            with gr.Accordion(label="GLB Extraction Settings", open=False):
                mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01)
                texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)

            with gr.Row():
                extract_glb_btn = gr.Button("Extract GLB", interactive=False)
                extract_gs_btn = gr.Button("Extract Gaussian", interactive=False)

            gr.Markdown("*NOTE: Gaussian file can be very large (~50MB), it will take a while to display and download.*")

        with gr.Column():
            video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
            model_output = gr.Model3D(label="Extracted GLB/Gaussian", height=300)

            with gr.Row():
                download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
                download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False)

    output_buf = gr.State()

    # Handlers
    demo.load(start_session)
    demo.unload(end_session)

    generate_btn.click(
        get_seed,
        inputs=[randomize_seed, seed],
        outputs=[seed],
    ).then(
        text_to_3d,
        inputs=[text_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
        outputs=[output_buf, video_output],
    ).then(
        lambda: tuple([gr.Button(interactive=True), gr.Button(interactive=True)]),
        outputs=[extract_glb_btn, extract_gs_btn],
    )

    video_output.clear(
        lambda: tuple([gr.Button(interactive=False), gr.Button(interactive=False)]),
        outputs=[extract_glb_btn, extract_gs_btn],
    )

    extract_glb_btn.click(
        extract_glb,
        inputs=[output_buf, mesh_simplify, texture_size],
        outputs=[model_output, download_glb],
    ).then(
        lambda: gr.Button(interactive=True),
        outputs=[download_glb],
    )

    extract_gs_btn.click(
        extract_gaussian,
        inputs=[output_buf],
        outputs=[model_output, download_gs],
    ).then(
        lambda: gr.Button(interactive=True),
        outputs=[download_gs],
    )

    model_output.clear(
        lambda: gr.Button(interactive=False),
        outputs=[download_glb],
    )


# -----------------------------
# Lanzamiento
# -----------------------------
if __name__ == "__main__":
    pipeline = TrellisTextTo3DPipeline.from_pretrained("cavargas10/TRELLIS-text-xlarge")
    pipeline.cuda()
    demo.launch()