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