| | """ |
| | TRELLIS.2 Text-to-3D Generator |
| | 🎨 Comic Classic Theme |
| | """ |
| |
|
| | import os |
| | import shutil |
| | import torch |
| | import numpy as np |
| | from PIL import Image |
| | import tempfile |
| | import uuid |
| | from typing import Tuple |
| | from datetime import datetime |
| |
|
| | import rerun as rr |
| | try: |
| | import rerun.blueprint as rrb |
| | except ImportError: |
| | rrb = None |
| |
|
| | from gradio_rerun import Rerun |
| | import gradio as gr |
| | from gradio_client import Client, handle_file |
| | import spaces |
| | from diffusers import ZImagePipeline |
| | from trellis2.pipelines import Trellis2ImageTo3DPipeline |
| | import o_voxel |
| |
|
| | os.environ["OPENCV_IO_ENABLE_OPENEXR"] = '1' |
| | os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" |
| | os.environ["ATTN_BACKEND"] = "flash_attn_3" |
| | os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'autotune_cache.json') |
| | os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1' |
| |
|
| | MAX_SEED = np.iinfo(np.int32).max |
| | TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp') |
| |
|
| | print("Loading Z-Image-Turbo...") |
| | try: |
| | z_pipe = ZImagePipeline.from_pretrained("Tongyi-MAI/Z-Image-Turbo", torch_dtype=torch.bfloat16, low_cpu_mem_usage=False) |
| | device = "cuda" if torch.cuda.is_available() else "cpu" |
| | z_pipe.to(device) |
| | except Exception as e: |
| | print(f"Failed to load Z-Image-Turbo: {e}") |
| | z_pipe = None |
| |
|
| | print("Loading TRELLIS.2...") |
| | try: |
| | trellis_pipeline = Trellis2ImageTo3DPipeline.from_pretrained('microsoft/TRELLIS.2-4B') |
| | trellis_pipeline.rembg_model = None |
| | trellis_pipeline.low_vram = False |
| | trellis_pipeline.cuda() |
| | except Exception as e: |
| | print(f"Failed to load TRELLIS.2: {e}") |
| | trellis_pipeline = None |
| |
|
| | rmbg_client = Client("briaai/BRIA-RMBG-2.0") |
| |
|
| |
|
| | 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)) |
| | if os.path.exists(user_dir): |
| | shutil.rmtree(user_dir) |
| |
|
| |
|
| | def remove_background(input: Image.Image) -> Image.Image: |
| | with tempfile.NamedTemporaryFile(suffix='.png') as f: |
| | input = input.convert('RGB') |
| | input.save(f.name) |
| | output = rmbg_client.predict(handle_file(f.name), api_name="/image")[0][0] |
| | output = Image.open(output) |
| | return output |
| |
|
| |
|
| | def preprocess_image(input: Image.Image) -> Image.Image: |
| | if input is None: |
| | return None |
| | has_alpha = False |
| | if input.mode == 'RGBA': |
| | alpha = np.array(input)[:, :, 3] |
| | if not np.all(alpha == 255): |
| | has_alpha = True |
| | max_size = max(input.size) |
| | scale = min(1, 1024 / max_size) |
| | if scale < 1: |
| | input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS) |
| | if has_alpha: |
| | output = input |
| | else: |
| | output = remove_background(input) |
| | output_np = np.array(output) |
| | alpha = output_np[:, :, 3] |
| | bbox = np.argwhere(alpha > 0.8 * 255) |
| | if bbox.size == 0: |
| | return output |
| | bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0]) |
| | center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2 |
| | size = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) |
| | size = int(size * 1) |
| | bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2 |
| | output = output.crop(bbox) |
| | output = np.array(output).astype(np.float32) / 255 |
| | output = output[:, :, :3] * output[:, :, 3:4] |
| | output = Image.fromarray((output * 255).astype(np.uint8)) |
| | return output |
| |
|
| |
|
| | def get_seed(randomize_seed: bool, seed: int) -> int: |
| | return np.random.randint(0, MAX_SEED) if randomize_seed else seed |
| |
|
| |
|
| | @spaces.GPU |
| | def generate_txt2img(prompt, progress=gr.Progress(track_tqdm=True)): |
| | if z_pipe is None: |
| | raise gr.Error("Z-Image-Turbo model failed to load.") |
| | if not prompt.strip(): |
| | raise gr.Error("Please enter a prompt.") |
| | device = "cuda" if torch.cuda.is_available() else "cpu" |
| | generator = torch.Generator(device).manual_seed(42) |
| | progress(0.1, desc="Generating Image...") |
| | try: |
| | result = z_pipe( |
| | prompt=prompt, |
| | negative_prompt=None, |
| | height=1024, |
| | width=1024, |
| | num_inference_steps=9, |
| | guidance_scale=0.0, |
| | generator=generator, |
| | ) |
| | return result.images[0] |
| | except Exception as e: |
| | raise gr.Error(f"Generation failed: {str(e)}") |
| |
|
| |
|
| | @spaces.GPU(duration=120) |
| | def generate_3d( |
| | image: Image.Image, seed: int, resolution: str, |
| | decimation_target: int, texture_size: int, |
| | ss_guidance_strength: float, ss_guidance_rescale: float, |
| | ss_sampling_steps: int, ss_rescale_t: float, |
| | shape_guidance: float, shape_rescale: float, |
| | shape_steps: int, shape_rescale_t: float, |
| | tex_guidance: float, tex_rescale: float, |
| | tex_steps: int, tex_rescale_t: float, |
| | req: gr.Request, progress=gr.Progress(track_tqdm=True) |
| | ) -> Tuple[str, str]: |
| | if image is None: |
| | raise gr.Error("Please provide an input image.") |
| | if trellis_pipeline is None: |
| | raise gr.Error("TRELLIS model is not loaded.") |
| | |
| | user_dir = os.path.join(TMP_DIR, str(req.session_hash)) |
| | os.makedirs(user_dir, exist_ok=True) |
| | |
| | progress(0.1, desc="Generating 3D...") |
| | try: |
| | outputs, latents = trellis_pipeline.run( |
| | image, seed=seed, preprocess_image=False, |
| | sparse_structure_sampler_params={"steps": ss_sampling_steps, "guidance_strength": ss_guidance_strength, "guidance_rescale": ss_guidance_rescale, "rescale_t": ss_rescale_t}, |
| | shape_slat_sampler_params={"steps": shape_steps, "guidance_strength": shape_guidance, "guidance_rescale": shape_rescale, "rescale_t": shape_rescale_t}, |
| | tex_slat_sampler_params={"steps": tex_steps, "guidance_strength": tex_guidance, "guidance_rescale": tex_rescale, "rescale_t": tex_rescale_t}, |
| | pipeline_type={"512": "512", "1024": "1024_cascade", "1536": "1536_cascade"}[resolution], |
| | return_latent=True, |
| | ) |
| | |
| | progress(0.7, desc="Processing Mesh...") |
| | mesh = outputs[0] |
| | mesh.simplify(1000000) |
| | |
| | progress(0.9, desc="Exporting GLB...") |
| | grid_size = latents[2] |
| | |
| | try: |
| | glb = o_voxel.postprocess.to_glb( |
| | vertices=mesh.vertices, faces=mesh.faces, attr_volume=mesh.attrs, |
| | coords=mesh.coords, attr_layout=trellis_pipeline.pbr_attr_layout, |
| | grid_size=grid_size, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], |
| | decimation_target=decimation_target, texture_size=texture_size, |
| | remesh=True, remesh_band=1, remesh_project=0, use_tqdm=True, |
| | ) |
| | except RuntimeError: |
| | glb = o_voxel.postprocess.to_glb( |
| | vertices=mesh.vertices, faces=mesh.faces, attr_volume=mesh.attrs, |
| | coords=mesh.coords, attr_layout=trellis_pipeline.pbr_attr_layout, |
| | grid_size=grid_size, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], |
| | decimation_target=decimation_target, texture_size=texture_size, |
| | remesh=False, remesh_band=1, remesh_project=0, use_tqdm=True, |
| | ) |
| | |
| | timestamp = datetime.now().strftime("%Y-%m-%dT%H%M%S") |
| | glb_path = os.path.join(user_dir, f'output_{timestamp}.glb') |
| | glb.export(glb_path, extension_webp=False) |
| | |
| | progress(0.95, desc="Creating Viewer...") |
| | run_id = str(uuid.uuid4()) |
| | rec = rr.new_recording(application_id="TRELLIS-3D-Viewer", recording_id=run_id) if hasattr(rr, "new_recording") else rr.RecordingStream(application_id="TRELLIS-3D-Viewer", recording_id=run_id) if hasattr(rr, "RecordingStream") else rr |
| | |
| | rec.log("world", rr.Clear(recursive=True), static=True) |
| | rec.log("world", rr.ViewCoordinates.RIGHT_HAND_Y_UP, static=True) |
| | rec.log("world/model", rr.Asset3D(path=glb_path), static=True) |
| | |
| | if rrb is not None: |
| | try: |
| | blueprint = rrb.Blueprint(rrb.Spatial3DView(origin="/world", name="3D View"), collapse_panels=True) |
| | rec.send_blueprint(blueprint) |
| | except: |
| | pass |
| | |
| | rrd_path = os.path.join(user_dir, f'output_{timestamp}.rrd') |
| | rec.save(rrd_path) |
| | torch.cuda.empty_cache() |
| | return rrd_path, glb_path |
| | |
| | except Exception as e: |
| | torch.cuda.empty_cache() |
| | raise gr.Error(f"Generation failed: {str(e)}") |
| |
|
| |
|
| | css = """ |
| | @import url('https://fonts.googleapis.com/css2?family=Bangers&family=Comic+Neue:wght@400;700&display=swap'); |
| | |
| | .gradio-container { |
| | background-color: #FEF9C3 !important; |
| | background-image: radial-gradient(#1F2937 1px, transparent 1px) !important; |
| | background-size: 20px 20px !important; |
| | min-height: 100vh !important; |
| | font-family: 'Comic Neue', cursive, sans-serif !important; |
| | } |
| | |
| | .huggingface-space-header, #space-header, .space-header, |
| | [class*="space-header"], .svelte-1ed2p3z, .space-header-badge, |
| | .header-badge, [data-testid="space-header"], .svelte-kqij2n, |
| | .svelte-1ax1toq, .embed-container > div:first-child { |
| | display: none !important; |
| | visibility: hidden !important; |
| | height: 0 !important; |
| | width: 0 !important; |
| | overflow: hidden !important; |
| | opacity: 0 !important; |
| | pointer-events: none !important; |
| | } |
| | |
| | footer, .footer, .gradio-container footer, .built-with, |
| | [class*="footer"], .gradio-footer, .main-footer, |
| | div[class*="footer"], .show-api, .built-with-gradio, |
| | a[href*="gradio.app"], a[href*="huggingface.co/spaces"] { |
| | display: none !important; |
| | visibility: hidden !important; |
| | height: 0 !important; |
| | padding: 0 !important; |
| | margin: 0 !important; |
| | } |
| | |
| | #col-container { max-width: 960px; margin: 0 auto; } |
| | |
| | .header-text h1 { |
| | font-family: 'Bangers', cursive !important; |
| | color: #1F2937 !important; |
| | font-size: 3.5rem !important; |
| | font-weight: 400 !important; |
| | text-align: center !important; |
| | margin-bottom: 0.5rem !important; |
| | text-shadow: 4px 4px 0px #FACC15, 6px 6px 0px #1F2937 !important; |
| | letter-spacing: 3px !important; |
| | -webkit-text-stroke: 2px #1F2937 !important; |
| | } |
| | |
| | .subtitle { |
| | text-align: center !important; |
| | font-family: 'Comic Neue', cursive !important; |
| | font-size: 1.2rem !important; |
| | color: #1F2937 !important; |
| | margin-bottom: 1.5rem !important; |
| | font-weight: 700 !important; |
| | } |
| | |
| | .gr-panel, .gr-box, .gr-form, .block, .gr-group { |
| | background: #FFFFFF !important; |
| | border: 3px solid #1F2937 !important; |
| | border-radius: 8px !important; |
| | box-shadow: 6px 6px 0px #1F2937 !important; |
| | transition: all 0.2s ease !important; |
| | } |
| | |
| | .gr-panel:hover, .block:hover { |
| | transform: translate(-2px, -2px) !important; |
| | box-shadow: 8px 8px 0px #1F2937 !important; |
| | } |
| | |
| | textarea, input[type="text"], input[type="number"] { |
| | background: #FFFFFF !important; |
| | border: 3px solid #1F2937 !important; |
| | border-radius: 8px !important; |
| | color: #1F2937 !important; |
| | font-family: 'Comic Neue', cursive !important; |
| | font-size: 1rem !important; |
| | font-weight: 700 !important; |
| | transition: all 0.2s ease !important; |
| | } |
| | |
| | textarea:focus, input[type="text"]:focus, input[type="number"]:focus { |
| | border-color: #3B82F6 !important; |
| | box-shadow: 4px 4px 0px #3B82F6 !important; |
| | outline: none !important; |
| | } |
| | |
| | .gr-button-primary, button.primary, .gr-button.primary { |
| | background: #3B82F6 !important; |
| | border: 3px solid #1F2937 !important; |
| | border-radius: 8px !important; |
| | color: #FFFFFF !important; |
| | font-family: 'Bangers', cursive !important; |
| | font-weight: 400 !important; |
| | font-size: 1.3rem !important; |
| | letter-spacing: 2px !important; |
| | padding: 14px 28px !important; |
| | box-shadow: 5px 5px 0px #1F2937 !important; |
| | transition: all 0.1s ease !important; |
| | text-shadow: 1px 1px 0px #1F2937 !important; |
| | } |
| | |
| | .gr-button-primary:hover, button.primary:hover, .gr-button.primary:hover { |
| | background: #2563EB !important; |
| | transform: translate(-2px, -2px) !important; |
| | box-shadow: 7px 7px 0px #1F2937 !important; |
| | } |
| | |
| | .gr-button-primary:active, button.primary:active, .gr-button.primary:active { |
| | transform: translate(3px, 3px) !important; |
| | box-shadow: 2px 2px 0px #1F2937 !important; |
| | } |
| | |
| | .gr-button-secondary, button.secondary { |
| | background: #EF4444 !important; |
| | border: 3px solid #1F2937 !important; |
| | border-radius: 8px !important; |
| | color: #FFFFFF !important; |
| | font-family: 'Bangers', cursive !important; |
| | font-weight: 400 !important; |
| | font-size: 1.1rem !important; |
| | letter-spacing: 1px !important; |
| | box-shadow: 4px 4px 0px #1F2937 !important; |
| | transition: all 0.1s ease !important; |
| | text-shadow: 1px 1px 0px #1F2937 !important; |
| | } |
| | |
| | .gr-button-secondary:hover, button.secondary:hover { |
| | background: #DC2626 !important; |
| | transform: translate(-2px, -2px) !important; |
| | box-shadow: 6px 6px 0px #1F2937 !important; |
| | } |
| | |
| | label, .gr-input-label, .gr-block-label { |
| | color: #1F2937 !important; |
| | font-family: 'Comic Neue', cursive !important; |
| | font-weight: 700 !important; |
| | font-size: 1rem !important; |
| | } |
| | |
| | .gr-file-upload { |
| | border: 3px dashed #1F2937 !important; |
| | border-radius: 8px !important; |
| | background: #FEF9C3 !important; |
| | } |
| | |
| | .gr-file-upload:hover { |
| | border-color: #3B82F6 !important; |
| | background: #EFF6FF !important; |
| | } |
| | |
| | ::-webkit-scrollbar { width: 12px; height: 12px; } |
| | ::-webkit-scrollbar-track { background: #FEF9C3; border: 2px solid #1F2937; } |
| | ::-webkit-scrollbar-thumb { background: #3B82F6; border: 2px solid #1F2937; border-radius: 0px; } |
| | ::-webkit-scrollbar-thumb:hover { background: #EF4444; } |
| | |
| | ::selection { background: #FACC15; color: #1F2937; } |
| | |
| | a { color: #3B82F6 !important; text-decoration: none !important; font-weight: 700 !important; } |
| | a:hover { color: #EF4444 !important; } |
| | |
| | @media (max-width: 768px) { |
| | .header-text h1 { |
| | font-size: 2.2rem !important; |
| | text-shadow: 3px 3px 0px #FACC15, 4px 4px 0px #1F2937 !important; |
| | } |
| | .gr-button-primary, button.primary { padding: 12px 20px !important; font-size: 1.1rem !important; } |
| | .gr-panel, .block { box-shadow: 4px 4px 0px #1F2937 !important; } |
| | } |
| | |
| | @media (prefers-color-scheme: dark) { |
| | .gradio-container { background-color: #FEF9C3 !important; } |
| | } |
| | """ |
| |
|
| | EXAMPLES_IMAGE = [f"example-images/A ({i}).webp" for i in range(1, 72)] |
| |
|
| | EXAMPLES_TEXT = [ |
| | "A Cat 3D model", "A realistic Cat 3D model", "A cartoon Cat 3D model", |
| | "A low poly Cat 3D", "A cyberpunk Cat 3D", "A robotic Cat 3D", |
| | "A Plane 3D model", "A fighter jet Plane 3D", "A vintage Plane 3D", |
| | "A Car 3D model", "A sports Car 3D", "A cyberpunk Car 3D", |
| | "A Shoe 3D model", "A sneaker Shoe 3D", "A boot Shoe 3D", |
| | "A Chair 3D model", "A Table 3D model", "A Robot 3D model", |
| | "A House 3D model", "A Spaceship 3D model", "A Motorcycle 3D model", |
| | ] |
| |
|
| | if __name__ == "__main__": |
| | os.makedirs(TMP_DIR, exist_ok=True) |
| |
|
| | with gr.Blocks(title="TRELLIS.2 Text-to-3D", delete_cache=(300, 300)) as demo: |
| | gr.LoginButton(value="Option: HuggingFace 'Login' for extra GPU quota +", size="sm") |
| | gr.HTML(f"<style>{css}</style>") |
| | |
| | gr.HTML(""" |
| | <div style="text-align: center; margin: 20px 0 10px 0;"> |
| | <a href="https://www.humangen.ai" target="_blank" style="text-decoration: none;"> |
| | <img src="https://img.shields.io/static/v1?label=🏠 HOME&message=HUMANGEN.AI&color=0000ff&labelColor=ffcc00&style=for-the-badge" alt="HOME"> |
| | </a> |
| | </div> |
| | """) |
| | |
| | gr.Markdown("# 🎮 TRELLIS.2 TEXT-TO-3D 🎮", elem_classes="header-text") |
| | gr.Markdown('<p class="subtitle">✨ Generate 3D models from text or images! 🚀</p>') |
| |
|
| | with gr.Row(): |
| | with gr.Column(scale=1, min_width=360): |
| | with gr.Tabs(): |
| | with gr.Tab("📝 Text-to-3D"): |
| | txt_prompt = gr.Textbox(label="💬 Prompt", placeholder="e.g. A Cat 3D model", lines=2) |
| | btn_gen_img = gr.Button("1️⃣ Generate Image", variant="primary") |
| | with gr.Tab("🖼️ Image-to-3D"): |
| | gr.Markdown("Upload an image directly.") |
| | |
| | image_prompt = gr.Image(label="📷 Input Image", format="png", image_mode="RGBA", type="pil", height=350) |
| |
|
| | with gr.Accordion(label="⚙️ 3D Settings", open=False): |
| | resolution = gr.Radio(["512", "1024", "1536"], label="Resolution", value="1024") |
| | seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1) |
| | randomize_seed = gr.Checkbox(label="🎲 Randomize Seed", value=True) |
| | |
| | decimation_target = gr.Slider(50000, 500000, label="Target Faces", value=150000, step=10000) |
| | texture_size = gr.Slider(512, 4096, label="Texture Size", value=1024, step=512) |
| | |
| | btn_gen_3d = gr.Button("2️⃣ Generate 3D", variant="primary") |
| |
|
| | with gr.Accordion(label="🔧 Advanced Sampler", open=False): |
| | gr.Markdown("**Stage 1: Sparse Structure**") |
| | ss_guidance_strength = gr.Slider(1.0, 10.0, value=7.5, label="Guidance") |
| | ss_guidance_rescale = gr.Slider(0.0, 1.0, value=0.7, label="Rescale") |
| | ss_sampling_steps = gr.Slider(1, 50, value=12, label="Steps") |
| | ss_rescale_t = gr.Slider(1.0, 6.0, value=5.0, label="Rescale T") |
| | |
| | gr.Markdown("**Stage 2: Shape**") |
| | shape_guidance = gr.Slider(1.0, 10.0, value=7.5, label="Guidance") |
| | shape_rescale = gr.Slider(0.0, 1.0, value=0.5, label="Rescale") |
| | shape_steps = gr.Slider(1, 50, value=12, label="Steps") |
| | shape_rescale_t = gr.Slider(1.0, 6.0, value=3.0, label="Rescale T") |
| | |
| | gr.Markdown("**Stage 3: Material**") |
| | tex_guidance = gr.Slider(1.0, 10.0, value=1.0, label="Guidance") |
| | tex_rescale = gr.Slider(0.0, 1.0, value=0.0, label="Rescale") |
| | tex_steps = gr.Slider(1, 50, value=12, label="Steps") |
| | tex_rescale_t = gr.Slider(1.0, 6.0, value=3.0, label="Rescale T") |
| |
|
| | with gr.Column(scale=2): |
| | gr.Markdown("### 🎯 3D Output") |
| | rerun_output = Rerun(label="3D Viewer", height=600) |
| | download_btn = gr.DownloadButton(label="3️⃣ Download GLB", variant="primary") |
| | |
| | gr.Examples(examples=[[img] for img in EXAMPLES_IMAGE], inputs=[image_prompt], label="🖼️ Image Examples") |
| | gr.Examples(examples=[[txt] for txt in EXAMPLES_TEXT], inputs=[txt_prompt], label="📝 Text Examples") |
| |
|
| | demo.load(start_session) |
| | demo.unload(end_session) |
| |
|
| | btn_gen_img.click(generate_txt2img, inputs=[txt_prompt], outputs=[image_prompt]).then( |
| | preprocess_image, inputs=[image_prompt], outputs=[image_prompt] |
| | ) |
| | |
| | image_prompt.upload(preprocess_image, inputs=[image_prompt], outputs=[image_prompt]) |
| |
|
| | btn_gen_3d.click(get_seed, inputs=[randomize_seed, seed], outputs=[seed]).then( |
| | generate_3d, |
| | inputs=[ |
| | image_prompt, seed, resolution, decimation_target, texture_size, |
| | ss_guidance_strength, ss_guidance_rescale, ss_sampling_steps, ss_rescale_t, |
| | shape_guidance, shape_rescale, shape_steps, shape_rescale_t, |
| | tex_guidance, tex_rescale, tex_steps, tex_rescale_t, |
| | ], |
| | outputs=[rerun_output, download_btn], |
| | ) |
| |
|
| | demo.launch() |
