Update app.py

app.py

@@ -1,3 +1,8 @@
+"""
+TRELLIS.2 Text-to-3D Generator
+🎨 Comic Classic Theme
+"""
+
 import os
 import shutil
 import torch
@@ -5,99 +10,16 @@ import numpy as np
 from PIL import Image
 import tempfile
 import uuid
-from typing import *
+from typing import Tuple
 from datetime import datetime
-from pathlib import Path
-
-from typing import Iterable
-from gradio.themes import Soft
-from gradio.themes.utils import colors, fonts, sizes
 
 import rerun as rr
-# Attempt to import blueprint for advanced view configuration
 try:
     import rerun.blueprint as rrb
 except ImportError:
     rrb = None
-    
-from gradio_rerun import Rerun
-
-# --- Theme Configuration ---
-colors.orange_red = colors.Color(
-    name="orange_red",
-    c50="#FFF0E5",
-    c100="#FFE0CC",
-    c200="#FFC299",
-    c300="#FFA366",
-    c400="#FF8533",
-    c500="#FF4500",
-    c600="#E63E00",
-    c700="#CC3700",
-    c800="#B33000",
-    c900="#992900",
-    c950="#802200",
-)
-
-class OrangeRedTheme(Soft):
-    def __init__(
-        self,
-        *,
-        primary_hue: colors.Color | str = colors.gray,
-        secondary_hue: colors.Color | str = colors.orange_red,
-        neutral_hue: colors.Color | str = colors.slate,
-        text_size: sizes.Size | str = sizes.text_lg,
-        font: fonts.Font | str | Iterable[fonts.Font | str] = (
-            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
-        ),
-        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
-            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
-        ),
-    ):
-        super().__init__(
-            primary_hue=primary_hue,
-            secondary_hue=secondary_hue,
-            neutral_hue=neutral_hue,
-            text_size=text_size,
-            font=font,
-            font_mono=font_mono,
-        )
-        super().set(
-            background_fill_primary="*primary_50",
-            background_fill_primary_dark="*primary_900",
-            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
-            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
-            button_primary_text_color="white",
-            button_primary_text_color_hover="white",
-            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
-            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
-            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
-            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
-            button_secondary_text_color="black",
-            button_secondary_text_color_hover="white",
-            button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)",
-            button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)",
-            button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)",
-            button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)",
-            slider_color="*secondary_500",
-            slider_color_dark="*secondary_600",
-            block_title_text_weight="600",
-            block_border_width="3px",
-            block_shadow="*shadow_drop_lg",
-            button_primary_shadow="*shadow_drop_lg",
-            button_large_padding="11px",
-            color_accent_soft="*primary_100",
-            block_label_background_fill="*primary_200",
-        )
-
-orange_red_theme = OrangeRedTheme()
-
-# --- Environment Setup ---
-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'
 
+from gradio_rerun import Rerun
 import gradio as gr
 from gradio_client import Client, handle_file
 import spaces
@@ -105,21 +27,20 @@ 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("Initializing models...")
-
 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,
-    )
+    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)
-    print("Z-Image-Turbo loaded.")
 except Exception as e:
     print(f"Failed to load Z-Image-Turbo: {e}")
     z_pipe = None
@@ -130,22 +51,24 @@ try:
     trellis_pipeline.rembg_model = None
     trellis_pipeline.low_vram = False
     trellis_pipeline.cuda()
-    print("TRELLIS.2 loaded.")
 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')
@@ -154,11 +77,10 @@ def remove_background(input: Image.Image) -> Image.Image:
         output = Image.open(output)
         return output
 
+
 def preprocess_image(input: Image.Image) -> Image.Image:
-    """Preprocess the input image: remove bg, crop, resize."""
     if input is None:
         return None
-        
     has_alpha = False
     if input.mode == 'RGBA':
         alpha = np.array(input)[:, :, 3]
@@ -172,12 +94,11 @@ def preprocess_image(input: Image.Image) -> Image.Image:
         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 
+        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])
@@ -189,21 +110,20 @@ def preprocess_image(input: Image.Image) -> Image.Image:
     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)):
-    """Generate Image using Z-Image Turbo"""
     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 Text-to-Image...")
+    progress(0.1, desc="Generating Image...")
     try:
         result = z_pipe(
             prompt=prompt,
@@ -216,176 +136,85 @@ def generate_txt2img(prompt, progress=gr.Progress(track_tqdm=True)):
         )
         return result.images[0]
     except Exception as e:
-        raise gr.Error(f"Z-Image Generation failed: {str(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),
+    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 Geometry...")
+    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,
-            },
+            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,
         )
         
-        # 2. Process Mesh
         progress(0.7, desc="Processing Mesh...")
         mesh = outputs[0]
-        # FIX: Aggressive simplification to prevent CUDA errors during post-processing
-        mesh.simplify(1000000) 
-        
-        # 3. Export to GLB
-        progress(0.9, desc="Baking Texture & Exporting GLB...")
+        mesh.simplify(1000000)
         
-        grid_size = latents[2] 
+        progress(0.9, desc="Exporting GLB...")
+        grid_size = latents[2]
         
-        # Safe export with fallback if high-quality hole filling fails
         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,
+                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 as e:
-            print(f"Warning: Post-processing failed with remesh=True. Error: {e}")
-            print("Retrying with remesh=False (Standard mesh generation)...")
+        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,
+                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,
             )
         
-        now = datetime.now()
-        timestamp = now.strftime("%Y-%m-%dT%H%M%S")
-        glb_path = os.path.join(user_dir, f'trellis_output_{timestamp}.glb')
-        
-        # FIX: extension_webp=False ensures compatibility with Rerun/Standard Viewers
+        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)
         
-        # --- Rerun Visualization Logic ---
-        progress(0.95, desc="Creating Rerun Visualization...")
-        
+        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
         
-        # Robustly handle different Rerun SDK versions
-        rec = None
-        if hasattr(rr, "new_recording"):
-            rec = rr.new_recording(application_id="TRELLIS-3D-Viewer", recording_id=run_id)
-        elif hasattr(rr, "RecordingStream"):
-            rec = rr.RecordingStream(application_id="TRELLIS-3D-Viewer", recording_id=run_id)
-        else:
-            rr.init("TRELLIS-3D-Viewer", recording_id=run_id, spawn=False)
-            rec = rr
-            
-        # 1. Clear State
         rec.log("world", rr.Clear(recursive=True), static=True)
-
-        # 2. Set View Coordinates: RIGHT_HAND_Y_UP
-        # This defines +Y as Up, +X as Right, and +Z as "Back" (towards the viewer).
-        # This effectively places the default camera in front of the object (at +Z).
         rec.log("world", rr.ViewCoordinates.RIGHT_HAND_Y_UP, static=True)
-
-        # 3. Add Axes Helpers (Red=X, Green=Y, Blue=Z)
-        try:
-            rec.log("world/axes/x", rr.Arrows3D(vectors=[[0.5, 0, 0]], colors=[[255, 0, 0]]), static=True)
-            rec.log("world/axes/y", rr.Arrows3D(vectors=[[0, 0.5, 0]], colors=[[0, 255, 0]]), static=True)
-            rec.log("world/axes/z", rr.Arrows3D(vectors=[[0, 0, 0.5]], colors=[[0, 0, 255]]), static=True)
-        except Exception:
-            pass
-
-        # 4. Log the 3D Model
         rec.log("world/model", rr.Asset3D(path=glb_path), static=True)
         
-        # 5. Send Blueprint (if supported) to force a clean 3D view
         if rrb is not None:
             try:
-                # Create a simple blueprint with a 3D view of "world"
-                blueprint = rrb.Blueprint(
-                    rrb.Spatial3DView(
-                        origin="/world",
-                        name="3D View",
-                    ),
-                    collapse_panels=True,
-                )
+                blueprint = rrb.Blueprint(rrb.Spatial3DView(origin="/world", name="3D View"), collapse_panels=True)
                 rec.send_blueprint(blueprint)
-            except Exception as e:
-                print(f"Blueprint creation failed (non-fatal): {e}")
-
-        # Save the Rerun recording (.rrd)
-        rrd_path = os.path.join(user_dir, f'trellis_output_{timestamp}.rrd')
+            except:
+                pass
+        
+        rrd_path = os.path.join(user_dir, f'output_{timestamp}.rrd')
         rec.save(rrd_path)
-
-        # Clean up
         torch.cuda.empty_cache()
         return rrd_path, glb_path
         
@@ -393,47 +222,235 @@ def generate_3d(
         torch.cuda.empty_cache()
         raise gr.Error(f"Generation failed: {str(e)}")
 
-css="""
-#col-container {
-    margin: 0 auto;
-    max-width: 960px;
+
+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; }
 }
-#main-title h1 {font-size: 2.4em !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(delete_cache=(300, 300)) as demo:
-        gr.Markdown("# **TRELLIS.2 (Text-to-3D)**", elem_id="main-title")
-        gr.Markdown("""
-        **Workflow:**
-        Generate a 3D asset directly by converting Text-to-Image → 3D or Image-to-3D, powered by [TRELLIS.2](https://huggingface.co/microsoft/TRELLIS.2-4B) and [Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo).
+    with gr.Blocks(title="TRELLIS.2 Text-to-3D", delete_cache=(300, 300)) as demo:
+        
+        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-Image-3D"):
-                        txt_prompt = gr.Textbox(label="Prompt", placeholder="eg. A Plane 3D", lines=2)
-                        btn_gen_img = gr.Button("1.Generate Image", variant="primary")
-                    with gr.Tab("Image-to-3D"):
-                        gr.Markdown("Upload an image directly if you have one.")
+                    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)
+                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="Generation Resolution", value="1024")
+                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)
-                    
+                    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", scale=2)
+                btn_gen_3d = gr.Button("2️⃣ Generate 3D", variant="primary")
 
-                with gr.Accordion(label="Advanced Sampler Settings", open=False):                
+                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")
@@ -453,193 +470,26 @@ if __name__ == "__main__":
                     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="Rerun 3D Viewer",
-                    height=600
-                )
-                
-                download_btn = gr.DownloadButton(label="3.Download GLB File", variant="primary")
-
-                gr.Examples(
-                    examples=[
-                        ["example-images/A (1).webp"],
-                        ["example-images/A (2).webp"],
-                        ["example-images/A (3).webp"],
-                        ["example-images/A (4).webp"],
-                        ["example-images/A (5).webp"],
-                        ["example-images/A (6).webp"],
-                        ["example-images/A (7).webp"],
-                        ["example-images/A (8).webp"],
-                        ["example-images/A (9).webp"],
-                        ["example-images/A (10).webp"],
-                        ["example-images/A (11).webp"],
-                        ["example-images/A (12).webp"],
-                        ["example-images/A (13).webp"],
-                        ["example-images/A (14).webp"],
-                        ["example-images/A (15).webp"],
-                        ["example-images/A (16).webp"],
-                        ["example-images/A (17).webp"],
-                        ["example-images/A (18).webp"],
-                        ["example-images/A (19).webp"],
-                        ["example-images/A (20).webp"],
-                        ["example-images/A (21).webp"],
-                        ["example-images/A (22).webp"],
-                        ["example-images/A (23).webp"],
-                        ["example-images/A (24).webp"],
-                        ["example-images/A (25).webp"],
-                        ["example-images/A (26).webp"],
-                        ["example-images/A (27).webp"],
-                        ["example-images/A (28).webp"],
-                        ["example-images/A (29).webp"],
-                        ["example-images/A (30).webp"],
-                        ["example-images/A (31).webp"],
-                        ["example-images/A (32).webp"],
-                        ["example-images/A (33).webp"],
-                        ["example-images/A (34).webp"],
-                        ["example-images/A (35).webp"],
-                        ["example-images/A (36).webp"],
-                        ["example-images/A (37).webp"],
-                        ["example-images/A (38).webp"],
-                        ["example-images/A (39).webp"],
-                        ["example-images/A (40).webp"],
-                        ["example-images/A (41).webp"],
-                        ["example-images/A (42).webp"],
-                        ["example-images/A (43).webp"],
-                        ["example-images/A (44).webp"],
-                        ["example-images/A (45).webp"],
-                        ["example-images/A (46).webp"],
-                        ["example-images/A (47).webp"],
-                        ["example-images/A (48).webp"],
-                        ["example-images/A (49).webp"],
-                        ["example-images/A (50).webp"],
-                        ["example-images/A (51).webp"],
-                        ["example-images/A (52).webp"],
-                        ["example-images/A (53).webp"],
-                        ["example-images/A (54).webp"],
-                        ["example-images/A (55).webp"],
-                        ["example-images/A (56).webp"],
-                        ["example-images/A (57).webp"],
-                        ["example-images/A (58).webp"],
-                        ["example-images/A (59).webp"],
-                        ["example-images/A (60).webp"],
-                        ["example-images/A (61).webp"],
-                        ["example-images/A (62).webp"],
-                        ["example-images/A (63).webp"],
-                        ["example-images/A (64).webp"],
-                        ["example-images/A (65).webp"],
-                        ["example-images/A (66).webp"],
-                        ["example-images/A (67).webp"],
-                        ["example-images/A (68).webp"],
-                        ["example-images/A (69).webp"],
-                        ["example-images/A (70).webp"],
-                        ["example-images/A (71).webp"],
-                    ],
-                    inputs=[image_prompt],
-                    label="Image Examples [image-to-3d]"
-                )
-                
-                gr.Examples(
-                    examples=[
-                        ["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 fluffy Cat 3D"],
-                        ["A fantasy Cat 3D creature"],
-                        ["A stylized Cat 3D"],
-                        ["A Cat 3D sculpture"],
-                
-                        ["A Plane 3D model"],
-                        ["A commercial Plane 3D"],
-                        ["A fighter jet Plane 3D"],
-                        ["A low poly Plane 3D"],
-                        ["A vintage Plane 3D"],
-                        ["A futuristic Plane 3D"],
-                        ["A cargo Plane 3D"],
-                        ["A private jet Plane 3D"],
-                        ["A toy Plane 3D"],
-                        ["A realistic Plane 3D"],
-                
-                        ["A Car 3D model"],
-                        ["A sports Car 3D"],
-                        ["A luxury Car 3D"],
-                        ["A low poly Car 3D"],
-                        ["A racing Car 3D"],
-                        ["A cyberpunk Car 3D"],
-                        ["A vintage Car 3D"],
-                        ["A futuristic Car 3D"],
-                        ["A SUV Car 3D"],
-                        ["A electric Car 3D"],
-                
-                        ["A Shoe 3D model"],
-                        ["A sneaker Shoe 3D"],
-                        ["A running Shoe 3D"],
-                        ["A leather Shoe 3D"],
-                        ["A high heel Shoe 3D"],
-                        ["A boot Shoe 3D"],
-                        ["A low poly Shoe 3D"],
-                        ["A futuristic Shoe 3D"],
-                        ["A sports Shoe 3D"],
-                        ["A casual Shoe 3D"],
-                
-                        ["A Chair 3D model"],
-                        ["A Table 3D model"],
-                        ["A Sofa 3D model"],
-                        ["A Lamp 3D model"],
-                        ["A Watch 3D model"],
-                        ["A Backpack 3D model"],
-                        ["A Drone 3D model"],
-                        ["A Robot 3D model"],
-                        ["A Smartphone 3D model"],
-                        ["A Headphones 3D model"],
-                
-                        ["A House 3D model"],
-                        ["A Skyscraper 3D model"],
-                        ["A Bridge 3D model"],
-                        ["A Castle 3D model"],
-                        ["A Spaceship 3D model"],
-                        ["A Rocket 3D model"],
-                        ["A Satellite 3D model"],
-                        ["A Tank 3D model"],
-                        ["A Motorcycle 3D model"],
-                        ["A Bicycle 3D model"]
-                    ],
-                    inputs=[txt_prompt],
-                    label="3D Prompt Examples [text-to-3d]"
-                )
+                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]
+        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],
-        )
+        image_prompt.upload(preprocess_image, inputs=[image_prompt], outputs=[image_prompt])
 
-        btn_gen_3d.click(
-            get_seed,
-            inputs=[randomize_seed, seed],
-            outputs=[seed],
-        ).then(
+        btn_gen_3d.click(get_seed, inputs=[randomize_seed, seed], outputs=[seed]).then(
             generate_3d,
             inputs=[
-                image_prompt, seed, resolution, 
-                decimation_target, texture_size,
+                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,
@@ -647,4 +497,4 @@ if __name__ == "__main__":
             outputs=[rerun_output, download_btn],
         )
 
-    demo.launch(theme=orange_red_theme, css=css, mcp_server=True, ssr_mode=False, show_error=True)
+    demo.launch()