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 logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - HF_SPACE_MULTIIMG - %(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) 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 al intentar eliminarlo. Es posible que ya haya sido limpiado.") def preprocess_images(images: List[Tuple[Image.Image, str]], req: gr.Request) -> List[Image.Image]: session_hash = str(req.session_hash) logging.info(f"[{session_hash}] Preprocesando {len(images)} imágenes.") images = [image[0] for image in images] processed_images = [pipeline.preprocess_image(image) for image in images] logging.info(f"[{session_hash}] Preprocesamiento completado.") return processed_images 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 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( multiimages: List[Tuple[Image.Image, str]], seed: int, ss_guidance_strength: float, ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int, multiimage_algo: Literal["multidiffusion", "stochastic"], req: gr.Request, ) -> Tuple[dict, str]: session_hash = str(req.session_hash) logging.info(f"[{session_hash}] Iniciando image_to_3d con {len(multiimages)} imágenes.") user_dir = os.path.join(TMP_DIR, session_hash) outputs = pipeline.run_multi_image( [image[0] for image in multiimages], 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, }, mode=multiimage_algo, ) 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 = [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 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 with gr.Blocks(delete_cache=(600, 600)) as demo: gr.Markdown(""" # UTPL - Conversión de Multiples Imágenes 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(): with gr.Column(): with gr.Tabs() as input_tabs: with gr.Tab(label="Multiple Images", id=1) as multiimage_input_tab: multiimage_prompt = gr.Gallery(label="Image Prompt", format="png", type="pil", height=300, columns=3) 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=12, 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=3.0, step=0.1) slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1) multiimage_algo = gr.Radio(["stochastic", "multidiffusion"], label="Multi-image Algorithm", value="stochastic") 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) extract_glb_btn = gr.Button("Extract GLB", interactive=False) with gr.Column(): video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300) model_output = gr.Model3D(label="Extracted GLB", height=300) download_glb = gr.DownloadButton(label="Download GLB", interactive=False) output_buf = gr.State() demo.load(start_session, inputs=None, outputs=None, api_name="start_session") demo.unload(end_session) multiimage_prompt.upload( preprocess_images, inputs=[multiimage_prompt], outputs=[multiimage_prompt], api_name="preprocess_images" ) generate_btn.click( get_seed, inputs=[randomize_seed, seed], outputs=[seed], api_name="get_seed" ).then( image_to_3d, inputs=[multiimage_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, multiimage_algo], outputs=[output_buf, video_output], api_name="image_to_3d" ).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], api_name="extract_glb" ).then( lambda: gr.Button(interactive=True), outputs=[download_glb], ) model_output.clear( lambda: gr.Button(interactive=False), outputs=[download_glb], ) # Lanzar la aplicación Gradio 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))) # Precargar rembg except: pass demo.launch(show_error=True)