app.py

import gradio as gr
import spaces 
import os
import shutil
os.environ['SPCONV_ALGO'] = 'native'
from typing import *
import torch
import numpy as np
import imageio
from easydict import EasyDict as edict
from PIL import Image
from trellis.pipelines import TrellisImageTo3DPipeline
from trellis.representations import Gaussian, MeshExtractResult
from trellis.utils import render_utils, postprocessing_utils
import pydantic
import logging

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

print(pydantic.__version__)

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)

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}")
    # Hacemos la eliminación más robusta.
    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 al intentar eliminarlo. Es posible que ya haya sido limpiado.")

def preprocess_image(image: Image.Image) -> Image.Image:
    logging.info("Preprocesando imagen...")
    processed_image = pipeline.preprocess_image(image)
    logging.info("Imagen preprocesada correctamente.")
    return processed_image

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, str]:
    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

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

@spaces.GPU
def image_to_3d(
    image: Image.Image,
    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 image_to_3d...")
    user_dir = os.path.join(TMP_DIR, session_hash)

    outputs = pipeline.run(
        image,
        seed=seed,
        formats=["gaussian", "mesh"],
        preprocess_image=False,
        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 del modelo 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 = 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 renderizado y estado empaquetado. 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 extraído. Devolviendo: {glb_path}")
    return glb_path, glb_path

def split_image(image: Image.Image) -> List[Image.Image]:
    image = np.array(image)
    alpha = image[..., 3]
    alpha = np.any(alpha>0, axis=0)
    start_pos = np.where(~alpha[:-1] & alpha[1:])[0].tolist()
    end_pos = np.where(alpha[:-1] & ~alpha[1:])[0].tolist()
    images = []
    for s, e in zip(start_pos, end_pos):
        images.append(Image.fromarray(image[:, s:e+1]))
    return [preprocess_image(image) for image in images]

with gr.Blocks(delete_cache=(600, 600)) as demo:
    gr.Markdown("""
    # UTPL - Conversión de Imágen 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/) (herramienta de código abierto para generación 3D)  
    **Propósito educativo:** Demostraciones académicas e Investigación en modelado 3D automático  
    """)
    
    with gr.Row(equal_height=False):
        # Left column (Controls)
        with gr.Column(scale=2, min_width=400):
            with gr.Tabs():
                with gr.Tab(label="Input Image"):
                    image_prompt = gr.Image(
                        label="Image Prompt", 
                        format="png", 
                        image_mode="RGBA", 
                        type="pil", 
                        height=300,
                        show_label=False
                    )
            
            with gr.Accordion(".Generation Settings", open=False):
                with gr.Column(variant="panel"):
                    seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
                    randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                    
                    with gr.Group():
                        gr.Markdown("#### Stage 1: Structure")
                        ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance", value=7.5, step=0.1)
                        ss_sampling_steps = gr.Slider(1, 50, label="Steps", value=12, step=1)
                    
                    with gr.Group():
                        gr.Markdown("#### Stage 2: Detail")
                        slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance", value=3.0, step=0.1)
                        slat_sampling_steps = gr.Slider(1, 50, label="Steps", value=12, step=1)
            
            generate_btn = gr.Button("Generate 3D Asset", variant="primary", size="lg")
            
            with gr.Accordion("GLB Export Settings", open=False):
                with gr.Column(variant="panel"):
                    mesh_simplify = gr.Slider(0.5, 0.98, label="Simplify Mesh", value=0.95, step=0.01)
                    texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
            
            extract_glb_btn = gr.Button("Export GLB", interactive=False, size="lg")

        with gr.Column(scale=3, min_width=600):
            with gr.Group():
                video_output = gr.Video(
                    label="3D Preview", 
                    autoplay=True, 
                    loop=True, 
                    height=300,
                    show_label=False
                )
                model_output = gr.Model3D(
                    label="3D Model Viewer", 
                    height=400
                )
            
            with gr.Row():
                download_glb = gr.DownloadButton(
                    label="Download GLB File", 
                    interactive=False,
                    variant="secondary",
                    size="lg"
                )

    output_buf = gr.State()

    demo.load(start_session)
    demo.unload(end_session)
    
    image_prompt.upload(
        preprocess_image,
        inputs=[image_prompt],
        outputs=[image_prompt],
    )

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

    video_output.clear(
        lambda: gr.Button(interactive=False),
        outputs=[extract_glb_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],
    )

    model_output.clear(
        lambda: gr.Button(interactive=False),
        outputs=[download_glb],
    )
    
if __name__ == "__main__":
    pipeline = TrellisImageTo3DPipeline.from_pretrained("cavargas10/TRELLIS")
    pipeline.cuda()
    try:
        pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
    except:
        pass
    demo.launch()