import gradio as gr import spaces from gradio_litmodel3d import LitModel3D 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 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): user_dir = os.path.join(TMP_DIR, str(req.session_hash)) os.makedirs(user_dir, exist_ok=True) def end_session(req: gr.Request): user_dir = os.path.join(TMP_DIR, str(req.session_hash)) shutil.rmtree(user_dir) def preprocess_image(image: Image.Image) -> Image.Image: processed_image = pipeline.preprocess_image(image) 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: return np.random.randint(0, MAX_SEED) if randomize_seed else 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]: user_dir = os.path.join(TMP_DIR, str(req.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, }, ) 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() 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]: user_dir = os.path.join(TMP_DIR, str(req.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() 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.9, 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") # Right column (Outputs) 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 = LitModel3D( label="3D Model Viewer", exposure=10.0, 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()