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 uuid from PIL import Image import logging from trellis.pipelines import TrellisImageTo3DPipeline from trellis.utils import 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) # Set up logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Initialize pipeline globally pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large") pipeline.cuda() try: pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))) # Preload rembg except: pass def start_session(req: gr.Request): user_dir = os.path.join(TMP_DIR, str(req.session_hash)) print(f'Creating user directory: {user_dir}') os.makedirs(user_dir, exist_ok=True) def end_session(req: gr.Request): user_dir = os.path.join(TMP_DIR, str(req.session_hash)) print(f'Removing user directory: {user_dir}') shutil.rmtree(user_dir) def preprocess_image(image: Image.Image, skip_if_rgba: bool = True) -> Image.Image: """ Preprocess the input image. Skip preprocessing if image is already RGBA and skip_if_rgba is True. """ if skip_if_rgba and image.mode == 'RGBA': return image return pipeline.preprocess_image(image) def get_seed(randomize_seed: bool, seed: int) -> int: """ Get the random seed. """ 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, mesh_simplify: float, # Now receiving the actual float value texture_size: int, # Now receiving the actual int value req: gr.Request, ) -> Tuple[dict, str]: """ Convert an image to a 3D model. Args: image (Image.Image): The input image. seed (int): The random seed. ss_guidance_strength (float): The guidance strength for sparse structure generation. ss_sampling_steps (int): The number of sampling steps for sparse structure generation. slat_guidance_strength (float): The guidance strength for structured latent generation. slat_sampling_steps (int): The number of sampling steps for structured latent generation. Returns: dict: The information of the generated 3D model. str: The path to the video of the 3D model. """ 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, }, ) # Direct GLB conversion trial_id = uuid.uuid4() glb_path = os.path.join(user_dir, f"{trial_id}.glb") try: torch.cuda.empty_cache() # Clear GPU memory before conversion glb = postprocessing_utils.to_glb( outputs['gaussian'][0], outputs['mesh'][0], simplify=mesh_simplify, texture_size=min(texture_size, 1024), # Limit texture size verbose=False, fill_holes=False, # Disable hole filling to save memory debug=False ) glb.export(glb_path) torch.cuda.empty_cache() # Clear GPU memory after conversion return glb_path except Exception as e: torch.cuda.empty_cache() # Ensure memory is cleared even if there's an error raise gr.Error(f"Failed to convert to GLB: {str(e)}") with gr.Blocks(delete_cache=(600, 600)) as demo: gr.Markdown(""" ## Image to 3D Asset with [TRELLIS](https://trellis3d.github.io/) * Upload an image and click "Generate" to create a 3D asset. If the image has alpha channel, it be used as the mask. Otherwise, we use `rembg` to remove the background. * If you find the generated 3D asset satisfactory, click "Extract GLB" to extract the GLB file and download it. """) with gr.Row(): with gr.Column(): image_prompt = gr.Image(label="Image Prompt", format="png", image_mode="RGBA", type="pil", height=300) 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) generate_btn = gr.Button("Generate 3D Model") 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.Column(): model_output = LitModel3D(label="Extracted GLB", exposure=20.0, height=300) download_glb = gr.DownloadButton(label="Download GLB", interactive=False) # Example images at the bottom of the page with gr.Row(): examples = gr.Examples( examples=[ f'assets/example_image/{image}' for image in os.listdir("assets/example_image") ], inputs=[image_prompt], fn=preprocess_image, outputs=[image_prompt], run_on_click=True, examples_per_page=64, ) # Handlers demo.load(start_session) demo.unload(end_session) image_prompt.upload( lambda img: preprocess_image(img, skip_if_rgba=False), 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, mesh_simplify, texture_size, ], outputs=[model_output], ).then( # This lambda function takes the glb_path and sets it as the download button's value lambda x: gr.Button(interactive=True, value=x), inputs=[model_output], outputs=[download_glb], ) model_output.clear( lambda: gr.Button(interactive=False), outputs=[download_glb], ) # Launch the Gradio app if __name__ == "__main__": demo.launch()