Update app.py

app.py

@@ -1,68 +1,61 @@
 import os
-import io
-import cv2
-import time
-import base64
 import shutil
-import tempfile
+import cv2
 import torch
 import numpy as np
+from PIL import Image
+import base64
+import io
+import tempfile
 import gradio as gr
 from gradio_client import Client, handle_file
+import spaces
 from pathlib import Path
-from typing import Tuple, List, Optional
-from PIL import Image
+from typing import *
 from datetime import datetime
 
-# --- Environment Configuration (Must be set before importing trellis2) ---
-os.environ["OPENCV_IO_ENABLE_OPENEXR"] = '1'
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-os.environ["ATTN_BACKEND"] = "flash_attn_3"
-# Adjust path if necessary or ensure autotune_cache.json exists in trellis2 dir
-try:
-    from trellis2 import modules
-    os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(modules.__file__)), 'autotune_cache.json')
-except:
-    pass
-os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1'
-
-import spaces # For GPU management in Hugging Face Spaces
-
-# --- Imports for Z-Image-Turbo ---
-from diffusers import ZImagePipeline
-
-# --- Imports for TRELLIS.2 ---
+# --- TRELLIS Specific Imports ---
+# Ensure these libraries are installed in your environment
 from trellis2.modules.sparse import SparseTensor
 from trellis2.pipelines import Trellis2ImageTo3DPipeline
 from trellis2.renderers import EnvMap
 from trellis2.utils import render_utils
 import o_voxel
 
+# --- Z-Image Specific Imports ---
+from diffusers import ZImagePipeline
+
 # ==========================================
-# Global Constants & CSS/JS
+# Configuration & 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"
+# Adjust path if necessary for your environment
+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')
 
-# Asset definitions
-MODES = [
-    {"name": "Normal", "icon": "assets/app/normal.png", "render_key": "normal"},
-    {"name": "Clay render", "icon": "assets/app/clay.png", "render_key": "clay"},
-    {"name": "Base color", "icon": "assets/app/basecolor.png", "render_key": "base_color"},
-    {"name": "HDRI forest", "icon": "assets/app/hdri_forest.png", "render_key": "shaded_forest"},
-    {"name": "HDRI sunset", "icon": "assets/app/hdri_sunset.png", "render_key": "shaded_sunset"},
-    {"name": "HDRI courtyard", "icon": "assets/app/hdri_courtyard.png", "render_key": "shaded_courtyard"},
-]
-STEPS = 8
-DEFAULT_MODE = 3
-DEFAULT_STEP = 3
+# ==========================================
+# CSS & JS Resources (From TRELLIS)
+# ==========================================
 
-CSS = """
-/* TRELLIS Custom Styles */
+css = """
+/* Overwrite Gradio Default Style */
 .stepper-wrapper { padding: 0; }
 .stepper-container { padding: 0; align-items: center; }
 .step-button { flex-direction: row; }
+.step-connector { transform: none; }
+.step-number { width: 16px; height: 16px; }
+.step-label { position: relative; bottom: 0; }
+.wrap.center.full { inset: 0; height: 100%; }
+.wrap.center.full.translucent { background: var(--block-background-fill); }
+.meta-text-center { display: block !important; position: absolute !important; top: unset !important; bottom: 0 !important; right: 0 !important; transform: unset !important; }
+
+/* Previewer */
 .previewer-container {
     position: relative;
     font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
@@ -76,53 +69,41 @@ CSS = """
     justify-content: center;
 }
 .previewer-container .tips-icon {
-    position: absolute; right: 10px; top: 10px; z-index: 10;
-    border-radius: 10px; color: #fff; background-color: var(--color-accent);
-    padding: 3px 6px; user-select: none;
+    position: absolute; right: 10px; top: 10px; z-index: 10; border-radius: 10px;
+    color: #fff; background-color: var(--color-accent); padding: 3px 6px; user-select: none;
 }
 .previewer-container .tips-text {
-    position: absolute; right: 10px; top: 50px; color: #fff;
-    background-color: var(--color-accent); border-radius: 10px;
-    padding: 6px; text-align: left; max-width: 300px; z-index: 10;
+    position: absolute; right: 10px; top: 50px; color: #fff; background-color: var(--color-accent);
+    border-radius: 10px; padding: 6px; text-align: left; max-width: 300px; z-index: 10;
     transition: all 0.3s; opacity: 0%; user-select: none;
 }
+.previewer-container .tips-text p { font-size: 14px; line-height: 1.2; }
 .tips-icon:hover + .tips-text { display: block; opacity: 100%; }
-.previewer-container .mode-row {
-    width: 100%; display: flex; gap: 8px; justify-content: center;
-    margin-bottom: 20px; flex-wrap: wrap;
-}
-.previewer-container .mode-btn {
-    width: 24px; height: 24px; border-radius: 50%; cursor: pointer;
-    opacity: 0.5; transition: all 0.2s; border: 2px solid #ddd; object-fit: cover;
-}
+
+/* Row 1: Display Modes */
+.previewer-container .mode-row { width: 100%; display: flex; gap: 8px; justify-content: center; margin-bottom: 20px; flex-wrap: wrap; }
+.previewer-container .mode-btn { width: 24px; height: 24px; border-radius: 50%; cursor: pointer; opacity: 0.5; transition: all 0.2s; border: 2px solid #ddd; object-fit: cover; }
 .previewer-container .mode-btn:hover { opacity: 0.9; transform: scale(1.1); }
 .previewer-container .mode-btn.active { opacity: 1; border-color: var(--color-accent); transform: scale(1.1); }
-.previewer-container .display-row {
-    margin-bottom: 20px; min-height: 400px; width: 100%; flex-grow: 1;
-    display: flex; justify-content: center; align-items: center;
-}
-.previewer-container .previewer-main-image {
-    max-width: 100%; max-height: 100%; flex-grow: 1; object-fit: contain; display: none;
-}
+
+/* Row 2: Display Image */
+.previewer-container .display-row { margin-bottom: 20px; min-height: 400px; width: 100%; flex-grow: 1; display: flex; justify-content: center; align-items: center; }
+.previewer-container .previewer-main-image { max-width: 100%; max-height: 100%; flex-grow: 1; object-fit: contain; display: none; }
 .previewer-container .previewer-main-image.visible { display: block; }
-.previewer-container .slider-row {
-    width: 100%; display: flex; flex-direction: column; align-items: center; gap: 10px; padding: 0 10px;
-}
-.previewer-container input[type=range] {
-    -webkit-appearance: none; width: 100%; max-width: 400px; background: transparent;
-}
-.previewer-container input[type=range]::-webkit-slider-runnable-track {
-    width: 100%; height: 8px; cursor: pointer; background: #ddd; border-radius: 5px;
-}
-.previewer-container input[type=range]::-webkit-slider-thumb {
-    height: 20px; width: 20px; border-radius: 50%; background: var(--color-accent);
-    cursor: pointer; -webkit-appearance: none; margin-top: -6px;
-    box-shadow: 0 2px 5px rgba(0,0,0,0.2); transition: transform 0.1s;
-}
+
+/* Row 3: Custom HTML Slider */
+.previewer-container .slider-row { width: 100%; display: flex; flex-direction: column; align-items: center; gap: 10px; padding: 0 10px; }
+.previewer-container input[type=range] { -webkit-appearance: none; width: 100%; max-width: 400px; background: transparent; }
+.previewer-container input[type=range]::-webkit-slider-runnable-track { width: 100%; height: 8px; cursor: pointer; background: #ddd; border-radius: 5px; }
+.previewer-container input[type=range]::-webkit-slider-thumb { height: 20px; width: 20px; border-radius: 50%; background: var(--color-accent); cursor: pointer; -webkit-appearance: none; margin-top: -6px; box-shadow: 0 2px 5px rgba(0,0,0,0.2); transition: transform 0.1s; }
+.previewer-container input[type=range]::-webkit-slider-thumb:hover { transform: scale(1.2); }
+
+/* Overwrite Previewer Block Style */
 .gradio-container .padded:has(.previewer-container) { padding: 0 !important; }
+.gradio-container:has(.previewer-container) [data-testid="block-label"] { position: absolute; top: 0; left: 0; }
 """
 
-HEAD = """
+head = """
 <script>
     function refreshView(mode, step) {
         const allImgs = document.querySelectorAll('.previewer-main-image');
@@ -140,6 +121,7 @@ HEAD = """
         const targetId = 'view-m' + mode + '-s' + step;
         const targetImg = document.getElementById(targetId);
         if (targetImg) targetImg.classList.add('visible');
+
         const allBtns = document.querySelectorAll('.mode-btn');
         allBtns.forEach((btn, idx) => {
             if (idx === mode) btn.classList.add('active');
@@ -151,68 +133,45 @@ HEAD = """
 </script>
 """
 
-EMPTY_HTML = f"""
+empty_html = f"""
 <div class="previewer-container">
-    <div style="opacity: 0.5; text-align: center;">
-        <p>3D Asset Preview will appear here.</p>
-    </div>
+    <svg style=" opacity: .5; height: var(--size-5); color: var(--body-text-color);"
+    xmlns="http://www.w3.org/2000/svg" width="100%" height="100%" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round" class="feather feather-image"><rect x="3" y="3" width="18" height="18" rx="2" ry="2"></rect><circle cx="8.5" cy="8.5" r="1.5"></circle><polyline points="21 15 16 10 5 21"></polyline></svg>
 </div>
 """
 
 # ==========================================
-# Model Loading
+# Model Loading & Global Variables
 # ==========================================
 
-print("Initializing models...")
-
-# 1. Load Z-Image-Turbo (Text to Image)
-print("Loading Z-Image-Turbo...")
-try:
-    t2i_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"
-    t2i_pipe.to(device)
-    print("Z-Image-Turbo loaded.")
-except Exception as e:
-    print(f"Failed to load Z-Image-Turbo: {e}")
-    t2i_pipe = None
-
-# 2. Load TRELLIS.2 (Image to 3D)
-print("Loading TRELLIS.2...")
-try:
-    pipeline_trellis = Trellis2ImageTo3DPipeline.from_pretrained('microsoft/TRELLIS.2-4B')
-    pipeline_trellis.rembg_model = None  # We use external Bria RMBG
-    pipeline_trellis.low_vram = False
-    pipeline_trellis.cuda()
-    print("TRELLIS.2 loaded.")
-except Exception as e:
-    print(f"Failed to load TRELLIS.2: {e}")
-    pipeline_trellis = None
-
-# 3. Load RMBG Client
-print("Loading RMBG Client...")
-try:
-    rmbg_client = Client("briaai/BRIA-RMBG-2.0")
-except Exception as e:
-    print(f"Failed to connect to RMBG client: {e}")
-    rmbg_client = None
-
-# 4. Load EnvMaps (Assuming assets folder exists)
+print("Loading Z-Image Turbo model...")
+# Text-to-Image Pipeline
+z_pipe = ZImagePipeline.from_pretrained(
+    "Tongyi-MAI/Z-Image-Turbo",
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=False,
+)
+if torch.cuda.is_available():
+    z_pipe.to("cuda")
+print("Z-Image Turbo loaded.")
+
+# Image-to-3D Pipeline Placeholders
+pipeline = None
+rmbg_client = None
 envmap = {}
-if os.path.exists('assets/hdri'):
-    try:
-        envmap = {
-            'forest': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread('assets/hdri/forest.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda')),
-            'sunset': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread('assets/hdri/sunset.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda')),
-            'courtyard': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread('assets/hdri/courtyard.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda')),
-        }
-    except Exception as e:
-        print(f"Warning: Could not load HDRIs: {e}")
-else:
-    print("Warning: 'assets/hdri' folder not found. Preview modes may fail.")
+
+# TRELLIS Settings
+MODES = [
+    {"name": "Normal", "icon": "assets/app/normal.png", "render_key": "normal"},
+    {"name": "Clay render", "icon": "assets/app/clay.png", "render_key": "clay"},
+    {"name": "Base color", "icon": "assets/app/basecolor.png", "render_key": "base_color"},
+    {"name": "HDRI forest", "icon": "assets/app/hdri_forest.png", "render_key": "shaded_forest"},
+    {"name": "HDRI sunset", "icon": "assets/app/hdri_sunset.png", "render_key": "shaded_sunset"},
+    {"name": "HDRI courtyard", "icon": "assets/app/hdri_courtyard.png", "render_key": "shaded_courtyard"},
+]
+STEPS = 8
+DEFAULT_MODE = 3
+DEFAULT_STEP = 3
 
 # ==========================================
 # Helper Functions
@@ -228,92 +187,51 @@ def image_to_base64(image):
 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):
-    try:
-        user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    if os.path.exists(user_dir):
         shutil.rmtree(user_dir)
-    except:
-        pass
 
 def remove_background(input: Image.Image) -> Image.Image:
-    """Removes background using Bria RMBG API via Gradio Client."""
-    if rmbg_client is None:
-        return input # Fallback
-    
-    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
+    with tempfile.NamedTemporaryFile(suffix='.png') as f:
         input = input.convert('RGB')
         input.save(f.name)
-        f_path = f.name
-    
-    try:
-        output_path = rmbg_client.predict(handle_file(f_path), api_name="/image")[0][0]
-        output = Image.open(output_path)
+        output = rmbg_client.predict(handle_file(f.name), api_name="/image")[0][0]
+        output = Image.open(output)
         return output
-    finally:
-        if os.path.exists(f_path):
-            os.remove(f_path)
 
 def preprocess_image(input: Image.Image) -> Image.Image:
-    """Preprocesses image (resizing, centering, BG removal) for TRELLIS."""
     if input is None:
         return None
-        
-    # Check Alpha
+    # if has alpha channel, use it directly; otherwise, remove background
     has_alpha = False
     if input.mode == 'RGBA':
         alpha = np.array(input)[:, :, 3]
         if not np.all(alpha == 255):
             has_alpha = True
-            
-    # Resize if too large
     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)
-        
-    # Remove BG if needed
     if has_alpha:
         output = input
     else:
         output = remove_background(input)
-        
-    # Centering and Cropping logic
     output_np = np.array(output)
-    # Ensure it has alpha now
-    if output_np.shape[2] == 4:
-        alpha = output_np[:, :, 3]
-        rows = np.any(alpha > 200, axis=1)
-        cols = np.any(alpha > 200, axis=0)
-        if np.any(rows) and np.any(cols): # Check if image is not empty
-            ymin, ymax = np.where(rows)[0][[0, -1]]
-            xmin, xmax = np.where(cols)[0][[0, -1]]
-            
-            w = xmax - xmin
-            h = ymax - ymin
-            size = max(w, h)
-            center_x, center_y = (xmin + xmax) / 2, (ymin + ymax) / 2
-            
-            # Add some padding
-            size = int(size * 1.1) 
-            
-            # Crop
-            left = max(0, int(center_x - size // 2))
-            top = max(0, int(center_y - size // 2))
-            right = min(output.width, int(center_x + size // 2))
-            bottom = min(output.height, int(center_y + size // 2))
-            
-            output = output.crop((left, top, right, bottom))
-            
-    # Premultiply alpha on black background logic for clean tensor conversion later? 
-    # Actually TRELLIS pipeline usually handles RGBA. 
-    # But let's standardize:
-    output_np = np.array(output).astype(np.float32) / 255
-    if output_np.shape[2] == 4:
-        # Premultiply
-        output_np[:, :, :3] * output_np[:, :, 3:4]
-    
-    output = Image.fromarray((output_np * 255).astype(np.uint8))
+    alpha = output_np[:, :, 3]
+    bbox = np.argwhere(alpha > 0.8 * 255)
+    if bbox.size == 0:
+        return output # Return original if no object found
+    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.1) # Added slight padding
+    bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2
+    output = output.crop(bbox)  # type: ignore
+    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 pack_state(latents: Tuple[SparseTensor, SparseTensor, int]) -> dict:
@@ -324,7 +242,7 @@ def pack_state(latents: Tuple[SparseTensor, SparseTensor, int]) -> dict:
         'coords': shape_slat.coords.cpu().numpy(),
         'res': res,
     }
-
+    
 def unpack_state(state: dict) -> Tuple[SparseTensor, SparseTensor, int]:
     shape_slat = SparseTensor(
         feats=torch.from_numpy(state['shape_slat_feats']).cuda(),
@@ -337,30 +255,37 @@ def get_seed(randomize_seed: bool, seed: int) -> int:
     return np.random.randint(0, MAX_SEED) if randomize_seed else seed
 
 # ==========================================
-# Main Processing Functions
+# Core Generation Functions (GPU)
 # ==========================================
 
 @spaces.GPU()
-def generate_text_to_image(prompt, progress=gr.Progress(track_tqdm=True)):
-    """Generates an image from text using Z-Image-Turbo."""
-    if t2i_pipe is None:
-        raise gr.Error("Text-to-Image Model not loaded.")
+def text_to_image(prompt, progress=gr.Progress(track_tqdm=True)):
+    """Generates image using Z-Image Turbo"""
+    if z_pipe is None:
+        raise gr.Error("Text-to-Image model failed to load.")
     if not prompt.strip():
-        raise gr.Error("Prompt is empty.")
-
+        raise gr.Error("Please enter a text prompt.")
+    
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    generator = torch.Generator(device).manual_seed(42) # Fixed seed for consistency demo, or make parametric
+    generator = torch.Generator(device).manual_seed(42) # Can be randomized if needed
     
-    result = t2i_pipe(
-        prompt=prompt,
-        negative_prompt=None,
-        height=1024,
-        width=1024,
-        num_inference_steps=9,
-        guidance_scale=0.0,
-        generator=generator,
-    )
-    return result.images[0]
+    progress(0.1, desc="Generating Text-to-Image...")
+    try:
+        result = z_pipe(
+            prompt=prompt,
+            negative_prompt=None,
+            height=1024,
+            width=1024,
+            num_inference_steps=9,
+            guidance_scale=0.0,
+            generator=generator,
+        )
+        image = result.images[0]
+        progress(1.0, desc="Image Generation Complete!")
+        # Automatically preprocess for 3D (remove bg if needed)
+        return preprocess_image(image)
+    except Exception as e:
+        raise gr.Error(f"Generation failed: {str(e)}")
 
 @spaces.GPU(duration=120)
 def image_to_3d(
@@ -381,24 +306,15 @@ def image_to_3d(
     tex_slat_rescale_t: float,
     req: gr.Request,
     progress=gr.Progress(track_tqdm=True),
-) -> Tuple[dict, str]:
-    
-    if pipeline_trellis is None:
-        raise gr.Error("TRELLIS Model not loaded.")
-
+) -> str:
     if image is None:
-        raise gr.Error("Input image is missing.")
-
-    # Ensure image is preprocessed (if it came directly from T2I, it has a background)
-    # If the user manually uploaded an RGBA, preprocess_image ensures it's clean.
-    # Note: Logic handled by calling preprocess_image in the Gradio event chain or inside here.
-    # We will assume the input 'image' to this function is the result of the Preprocess step.
+        raise gr.Error("Please upload or generate an image first.")
 
     # --- Sampling ---
-    outputs, latents = pipeline_trellis.run(
+    outputs, latents = pipeline.run(
         image,
         seed=seed,
-        preprocess_image=False, # We assume input is already preprocessed
+        preprocess_image=False, # We preprocessed before this step
         sparse_structure_sampler_params={
             "steps": ss_sampling_steps,
             "guidance_strength": ss_guidance_strength,
@@ -424,64 +340,40 @@ def image_to_3d(
         }[resolution],
         return_latent=True,
     )
-    
     mesh = outputs[0]
-    mesh.simplify(16777216) 
-    
-    # Render Preview
-    images_rendered = render_utils.render_snapshot(mesh, resolution=1024, r=2, fov=36, nviews=STEPS, envmap=envmap)
+    mesh.simplify(16777216) # nvdiffrast limit
+    images = render_utils.render_snapshot(mesh, resolution=1024, r=2, fov=36, nviews=STEPS, envmap=envmap)
     state = pack_state(latents)
     torch.cuda.empty_cache()
     
     # --- HTML Construction ---
     images_html = ""
     for m_idx, mode in enumerate(MODES):
-        if mode['render_key'] not in images_rendered: continue # skip if missing hdri
         for s_idx in range(STEPS):
             unique_id = f"view-m{m_idx}-s{s_idx}"
             is_visible = (m_idx == DEFAULT_MODE and s_idx == DEFAULT_STEP)
             vis_class = "visible" if is_visible else ""
-            img_base64 = image_to_base64(Image.fromarray(images_rendered[mode['render_key']][s_idx]))
-            
-            images_html += f"""
-                <img id="{unique_id}" 
-                     class="previewer-main-image {vis_class}" 
-                     src="{img_base64}" 
-                     loading="eager">
-            """
+            img_base64 = image_to_base64(Image.fromarray(images[mode['render_key']][s_idx]))
+            images_html += f"""<img id="{unique_id}" class="previewer-main-image {vis_class}" src="{img_base64}" loading="eager">"""
     
     btns_html = ""
-    for idx, mode in enumerate(MODES): 
-        if mode['render_key'] not in images_rendered: continue
+    for idx, mode in enumerate(MODES):        
         active_class = "active" if idx == DEFAULT_MODE else ""
-        btns_html += f"""
-            <img src="{mode['icon_base64']}" 
-                 class="mode-btn {active_class}" 
-                 onclick="selectMode({idx})"
-                 title="{mode['name']}">
-        """
+        btns_html += f"""<img src="{mode['icon_base64']}" class="mode-btn {active_class}" onclick="selectMode({idx})" title="{mode['name']}">"""
     
     full_html = f"""
     <div class="previewer-container">
         <div class="tips-wrapper">
             <div class="tips-icon">💡Tips</div>
-            <div class="tips-text">
-                <p>● <b>Render Mode</b> - Click buttons to switch render modes.</p>
-                <p>● <b>View Angle</b> - Drag slider to rotate.</p>
-            </div>
-        </div>
-        <div class="display-row">
-            {images_html}
-        </div>
-        <div class="mode-row" id="btn-group">
-            {btns_html}
+            <div class="tips-text"><p>● <b>Render Mode</b> - Switch render modes.</p><p>● <b>View Angle</b> - Drag slider.</p></div>
         </div>
+        <div class="display-row">{images_html}</div>
+        <div class="mode-row" id="btn-group">{btns_html}</div>
         <div class="slider-row">
             <input type="range" id="custom-slider" min="0" max="{STEPS - 1}" value="{DEFAULT_STEP}" step="1" oninput="onSliderChange(this.value)">
         </div>
     </div>
     """
-    
     return state, full_html
 
 @spaces.GPU(duration=120)
@@ -492,18 +384,19 @@ def extract_glb(
     req: gr.Request,
     progress=gr.Progress(track_tqdm=True),
 ) -> Tuple[str, str]:
-    
+    if state is None:
+        raise gr.Error("No 3D model generated yet.")
+        
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     shape_slat, tex_slat, res = unpack_state(state)
-    mesh = pipeline_trellis.decode_latent(shape_slat, tex_slat, res)[0]
+    mesh = pipeline.decode_latent(shape_slat, tex_slat, res)[0]
     mesh.simplify(16777216)
-    
     glb = o_voxel.postprocess.to_glb(
         vertices=mesh.vertices,
         faces=mesh.faces,
         attr_volume=mesh.attrs,
         coords=mesh.coords,
-        attr_layout=pipeline_trellis.pbr_attr_layout,
+        attr_layout=pipeline.pbr_attr_layout,
         grid_size=res,
         aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
         decimation_target=decimation_target,
@@ -513,7 +406,6 @@ def extract_glb(
         remesh_project=0,
         use_tqdm=True,
     )
-    
     now = datetime.now()
     timestamp = now.strftime("%Y-%m-%dT%H%M%S") + f".{now.microsecond // 1000:03d}"
     os.makedirs(user_dir, exist_ok=True)
@@ -523,135 +415,147 @@ def extract_glb(
     return glb_path, glb_path
 
 # ==========================================
-# Gradio Interface
+# Gradio UI Blocks
 # ==========================================
 
-if __name__ == "__main__":
-    os.makedirs(TMP_DIR, exist_ok=True)
+with gr.Blocks(delete_cache=(600, 600), title="TRELLIS.2-3D") as demo:
+    gr.Markdown("""
+    # TRELLIS.2-3D: Text to 3D Generation
+    **Model 1:** [Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) (Text to Image)  
+    **Model 2:** [TRELLIS.2](https://microsoft.github.io/TRELLIS.2) (Image to 3D)
     
-    # Pre-calculate base64 for icons to avoid FS lag
-    for i in range(len(MODES)):
-        if os.path.exists(MODES[i]['icon']):
-            icon = Image.open(MODES[i]['icon'])
-            MODES[i]['icon_base64'] = image_to_base64(icon)
-        else:
-            MODES[i]['icon_base64'] = "" # Handle missing assets
-
-    with gr.Blocks(css=CSS, head=HEAD, delete_cache=(600, 600)) as demo:
-        gr.Markdown("""
-        # TRELLIS.2-3D
-        ### Text-to-Image (Z-Image-Turbo) + Image-to-3D (TRELLIS.2) Pipeline
-        """)
-        
-        with gr.Row():
-            # Left Column: Inputs & Text-to-Image
-            with gr.Column(scale=1):
-                gr.Markdown("### 1. Generate Image")
-                text_prompt = gr.Textbox(label="Text Prompt", placeholder="A 3D rendering of a cute isometric house...")
-                gen_image_btn = gr.Button("Generate Image from Text", variant="secondary")
+    *Generate an image from text, then convert it to a 3D Asset.*
+    """)
+    
+    with gr.Row():
+        # --- Column 1: Controls ---
+        with gr.Column(scale=1, min_width=360):
+            gr.Markdown("### 1. Generate Image (Z-Image)")
+            text_prompt = gr.Textbox(label="Text Prompt", placeholder="A cute isometric 3D render of a robot...")
+            gen_image_btn = gr.Button("Generate Image", variant="primary")
+            
+            gr.Markdown("### 2. Configure 3D (TRELLIS)")
+            image_input = gr.Image(label="Input Image (Auto-filled or Upload)", format="png", image_mode="RGBA", type="pil", height=300)
+            
+            with gr.Accordion("3D Generation Settings", open=True):
+                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)
+                
+            with gr.Accordion(label="Advanced 3D Settings", open=False):                
+                gr.Markdown("Stage 1: Sparse Structure")
+                ss_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance", value=7.5)
+                ss_guidance_rescale = gr.Slider(0.0, 1.0, label="Rescale", value=0.7)
+                ss_sampling_steps = gr.Slider(1, 50, label="Steps", value=12)
+                ss_rescale_t = gr.Slider(1.0, 6.0, label="Rescale T", value=5.0)
                 
-                gr.Markdown("### 2. Prepare for 3D")
-                # This Image component acts as the bridge. 
-                # It accepts output from T2I OR user upload.
-                image_input = gr.Image(label="Input Image (Auto-Preprocessed)", type="pil", image_mode="RGBA", height=300)
+                gr.Markdown("Stage 2: Shape")
+                shape_slat_guidance_strength = gr.Slider(1.0, 10.0, value=7.5)
+                shape_slat_guidance_rescale = gr.Slider(0.0, 1.0, value=0.5)
+                shape_slat_sampling_steps = gr.Slider(1, 50, value=12)
+                shape_slat_rescale_t = gr.Slider(1.0, 6.0, value=3.0)
                 
-                with gr.Accordion("Image-to-3D Settings", open=True):
-                    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)
+                gr.Markdown("Stage 3: Material")
+                tex_slat_guidance_strength = gr.Slider(1.0, 10.0, value=1.0)
+                tex_slat_guidance_rescale = gr.Slider(0.0, 1.0, value=0.0)
+                tex_slat_sampling_steps = gr.Slider(1, 50, value=12)
+                tex_slat_rescale_t = gr.Slider(1.0, 6.0, value=3.0) 
+
+            generate_3d_btn = gr.Button("Generate 3D Asset", variant="primary")
+
+        # --- Column 2: Output ---
+        with gr.Column(scale=2):
+            with gr.Walkthrough(selected=0) as walkthrough:
+                with gr.Step("Preview", id=0):
+                    preview_output = gr.HTML(empty_html, label="3D Asset Preview", show_label=True, container=True)
                     
-                    with gr.Accordion("Advanced Parameters", open=False):
-                        gr.Markdown("**Stage 1: Sparse Structure**")
-                        ss_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance", value=7.5, step=0.1)
-                        ss_guidance_rescale = gr.Slider(0.0, 1.0, label="Rescale", value=0.7, step=0.01)
-                        ss_sampling_steps = gr.Slider(1, 50, label="Steps", value=12, step=1)
-                        ss_rescale_t = gr.Slider(1.0, 6.0, label="Rescale T", value=5.0, step=0.1)
-                        
-                        gr.Markdown("**Stage 2: Shape**")
-                        shape_slat_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance", value=7.5, step=0.1)
-                        shape_slat_guidance_rescale = gr.Slider(0.0, 1.0, label="Rescale", value=0.5, step=0.01)
-                        shape_slat_sampling_steps = gr.Slider(1, 50, label="Steps", value=12, step=1)
-                        shape_slat_rescale_t = gr.Slider(1.0, 6.0, label="Rescale T", value=3.0, step=0.1)
-                        
-                        gr.Markdown("**Stage 3: Texture**")
-                        tex_slat_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance", value=1.0, step=0.1)
-                        tex_slat_guidance_rescale = gr.Slider(0.0, 1.0, label="Rescale", value=0.0, step=0.01)
-                        tex_slat_sampling_steps = gr.Slider(1, 50, label="Steps", value=12, step=1)
-                        tex_slat_rescale_t = gr.Slider(1.0, 6.0, label="Rescale T", value=3.0, step=0.1)
-
-                gen_3d_btn = gr.Button("Generate 3D Model", variant="primary")
-
-            # Right Column: 3D Preview & Export
-            with gr.Column(scale=2):
-                gr.Markdown("### 3. 3D Preview")
-                # We use a Walkthrough to switch between Preview HTML and GLB Viewer
-                with gr.Walkthrough(selected=0) as walkthrough:
-                    with gr.Step("Preview", id=0):
-                        preview_output = gr.HTML(EMPTY_HTML, label="3D Preview", container=True)
-                        gr.Markdown("*(If the preview is black, verify 'assets/hdri' files exist)*")
-                        
-                        with gr.Row():
-                            decimation_target = gr.Slider(100000, 500000, label="Mesh Decimation", value=300000, step=10000)
-                            texture_size = gr.Slider(1024, 4096, label="Texture Size", value=2048, step=1024)
-                        extract_btn = gr.Button("Extract & Download GLB")
-                        
-                    with gr.Step("Result", id=1):
-                        glb_output = gr.Model3D(label="Extracted GLB", height=600, display_mode="solid", clear_color=(0.2, 0.2, 0.2, 1.0))
-                        download_btn = gr.DownloadButton(label="Download .glb File")
-                        back_btn = gr.Button("Back to Preview")
-
-        # Hidden State to store TRELLIS latent representation
-        output_state = gr.State()
-
-        # ====================
-        # Event Handling
-        # ====================
-        
-        demo.load(start_session)
-        demo.unload(end_session)
-
-        # 1. Text to Image
-        gen_image_btn.click(
-            generate_text_to_image,
-            inputs=[text_prompt],
-            outputs=[image_input]
-        )
+                    gr.Markdown("### 3. Extract & Download")
+                    with gr.Row():
+                        decimation_target = gr.Slider(100000, 500000, label="Decimation Target", value=300000, step=10000)
+                        texture_size = gr.Slider(1024, 4096, label="Texture Size", value=2048, step=1024)
+                    extract_btn = gr.Button("Extract GLB")
+                    
+                with gr.Step("Extract", id=1):
+                    glb_output = gr.Model3D(label="Extracted GLB", height=600, show_label=True, display_mode="solid", clear_color=(0.25, 0.25, 0.25, 1.0))
+                    download_btn = gr.DownloadButton(label="Download GLB")
+                    
+    output_buf = gr.State()
+    
+    # --- Event Handlers ---
+    
+    demo.load(start_session)
+    demo.unload(end_session)
+    
+    # 1. Text to Image
+    gen_image_btn.click(
+        text_to_image,
+        inputs=[text_prompt],
+        outputs=[image_input]
+    )
+    
+    # Handle manual upload preprocessing
+    image_input.upload(
+        preprocess_image,
+        inputs=[image_input],
+        outputs=[image_input],
+    )
 
-        # 2. Image Preprocessing (Auto-trigger when image changes)
-        image_input.change(
-            preprocess_image,
-            inputs=[image_input],
-            outputs=[image_input]
-        )
+    # 2. Image to 3D
+    generate_3d_btn.click(
+        get_seed,
+        inputs=[randomize_seed, seed],
+        outputs=[seed],
+    ).then(
+        lambda: gr.Walkthrough(selected=0), outputs=walkthrough
+    ).then(
+        image_to_3d,
+        inputs=[
+            image_input, seed, resolution,
+            ss_guidance_strength, ss_guidance_rescale, ss_sampling_steps, ss_rescale_t,
+            shape_slat_guidance_strength, shape_slat_guidance_rescale, shape_slat_sampling_steps, shape_slat_rescale_t,
+            tex_slat_guidance_strength, tex_slat_guidance_rescale, tex_slat_sampling_steps, tex_slat_rescale_t,
+        ],
+        outputs=[output_buf, preview_output],
+    )
+    
+    # 3. Extract GLB
+    extract_btn.click(
+        lambda: gr.Walkthrough(selected=1), outputs=walkthrough
+    ).then(
+        extract_glb,
+        inputs=[output_buf, decimation_target, texture_size],
+        outputs=[glb_output, download_btn],
+    )
 
-        # 3. Image to 3D Generation
-        gen_3d_btn.click(
-            get_seed,
-            inputs=[randomize_seed, seed],
-            outputs=[seed],
-        ).then(
-            lambda: gr.Walkthrough(selected=0), outputs=walkthrough
-        ).then(
-            image_to_3d,
-            inputs=[
-                image_input, seed, resolution,
-                ss_guidance_strength, ss_guidance_rescale, ss_sampling_steps, ss_rescale_t,
-                shape_slat_guidance_strength, shape_slat_guidance_rescale, shape_slat_sampling_steps, shape_slat_rescale_t,
-                tex_slat_guidance_strength, tex_slat_guidance_rescale, tex_slat_sampling_steps, tex_slat_rescale_t,
-            ],
-            outputs=[output_state, preview_output],
-        )
+# ==========================================
+# Main Execution
+# ==========================================
 
-        # 4. Extract GLB
-        extract_btn.click(
-            lambda: gr.Walkthrough(selected=1), outputs=walkthrough
-        ).then(
-            extract_glb,
-            inputs=[output_state, decimation_target, texture_size],
-            outputs=[glb_output, download_btn],
-        )
-        
-        # 5. Back button
-        back_btn.click(lambda: gr.Walkthrough(selected=0), outputs=walkthrough)
+if __name__ == "__main__":
+    os.makedirs(TMP_DIR, exist_ok=True)
 
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    # Initialize TRELLIS dependencies and assets
+    try:
+        # Load Icons
+        for i in range(len(MODES)):
+            if os.path.exists(MODES[i]['icon']):
+                icon = Image.open(MODES[i]['icon'])
+                MODES[i]['icon_base64'] = image_to_base64(icon)
+            else:
+                print(f"Warning: Icon not found at {MODES[i]['icon']}")
+                MODES[i]['icon_base64'] = "" # Fallback or placeholder
+
+        # Initialize RMBG and TRELLIS
+        print("Loading RMBG and TRELLIS models...")
+        rmbg_client = Client("briaai/BRIA-RMBG-2.0")
+        pipeline = Trellis2ImageTo3DPipeline.from_pretrained('microsoft/TRELLIS.2-4B')
+        pipeline.rembg_model = None
+        pipeline.low_vram = False
+        pipeline.cuda()
+        
+        # Load Env Maps
+        envmap_files = {
+            'forest': 'assets/hdri/forest.exr',
+            'sunset': 'assets/hdri/sunset.exr',
+            'courtyard': 'assets/hdri/courtyard.exr'
+        }