Update app.py

app.py

@@ -1,31 +1,37 @@
 import os
-import shutil
+import io
 import cv2
+import time
+import base64
+import shutil
+import tempfile
 import torch
 import numpy as np
+import gradio as gr
+from gradio_client import Client, handle_file
+from pathlib import Path
+from typing import Tuple, List, Optional
 from PIL import Image
-import base64
-import io
-import tempfile
-from typing import *
 from datetime import datetime
-from pathlib import Path
 
-# --- Environment Setup (Must be before other imports) ---
+# --- 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"
-# Set autotune cache relative to this file
-os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'autotune_cache.json')
+# 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'
 
-# --- Third Party Imports ---
-import gradio as gr
-from gradio_client import Client, handle_file
-import spaces
+import spaces # For GPU management in Hugging Face Spaces
+
+# --- Imports for Z-Image-Turbo ---
 from diffusers import ZImagePipeline
 
-# --- TRELLIS Specific Imports ---
+# --- Imports for TRELLIS.2 ---
 from trellis2.modules.sparse import SparseTensor
 from trellis2.pipelines import Trellis2ImageTo3DPipeline
 from trellis2.renderers import EnvMap
@@ -33,13 +39,13 @@ from trellis2.utils import render_utils
 import o_voxel
 
 # ==========================================
-# Global Configuration & Assets
+# Global Constants & CSS/JS
 # ==========================================
 
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 
-# TRELLIS Render Modes
+# Asset definitions
 MODES = [
     {"name": "Normal", "icon": "assets/app/normal.png", "render_key": "normal"},
     {"name": "Clay render", "icon": "assets/app/clay.png", "render_key": "clay"},
@@ -52,57 +58,71 @@ STEPS = 8
 DEFAULT_MODE = 3
 DEFAULT_STEP = 3
 
-# ==========================================
-# CSS & JavaScript (For Custom Previewer)
-# ==========================================
-
-css = """
-/* Overwrite Gradio Default Style */
+CSS = """
+/* TRELLIS Custom Styles */
 .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;
-    width: 100%; height: 722px; margin: 0 auto; padding: 20px;
-    display: flex; flex-direction: column; align-items: center; justify-content: center;
+    position: relative;
+    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
+    width: 100%;
+    height: 722px;
+    margin: 0 auto;
+    padding: 20px;
+    display: flex;
+    flex-direction: column;
+    align-items: center;
+    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;
+}
+.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;
+    transition: all 0.3s; opacity: 0%; user-select: none;
 }
-.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; }
-.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; 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%; }
-
-/* 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-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); }
-
-/* 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 .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; }
-
-/* 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 */
+.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;
+}
 .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');
@@ -120,7 +140,6 @@ 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');
@@ -132,10 +151,11 @@ head = """
 </script>
 """
 
-empty_html = f"""
+EMPTY_HTML = f"""
 <div class="previewer-container">
-    <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 style="opacity: 0.5; text-align: center;">
+        <p>3D Asset Preview will appear here.</p>
+    </div>
 </div>
 """
 
@@ -145,43 +165,54 @@ empty_html = f"""
 
 print("Initializing models...")
 
-# 1. Z-Image-Turbo (Text to Image)
+# 1. Load Z-Image-Turbo (Text to Image)
 print("Loading Z-Image-Turbo...")
 try:
-    z_pipe = ZImagePipeline.from_pretrained(
+    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"
-    z_pipe.to(device)
+    t2i_pipe.to(device)
     print("Z-Image-Turbo loaded.")
 except Exception as e:
     print(f"Failed to load Z-Image-Turbo: {e}")
-    z_pipe = None
+    t2i_pipe = None
 
-# 2. TRELLIS.2 (Image to 3D)
+# 2. Load TRELLIS.2 (Image to 3D)
 print("Loading TRELLIS.2...")
-# Initialize on startup
-trellis_pipeline = Trellis2ImageTo3DPipeline.from_pretrained('microsoft/TRELLIS.2-4B')
-trellis_pipeline.rembg_model = None
-trellis_pipeline.low_vram = False
-trellis_pipeline.cuda()
-
-# 3. Background Remover
-rmbg_client = Client("briaai/BRIA-RMBG-2.0")
+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
 
-# 4. HDRI Maps for TRELLIS
-envmap = {}
-# Try to load assets, handle gracefully if running in a basic environment
+# 3. Load RMBG Client
+print("Loading RMBG Client...")
 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')),
-    }
+    rmbg_client = Client("briaai/BRIA-RMBG-2.0")
 except Exception as e:
-    print(f"Warning: Could not load HDRI assets. {e}")
+    print(f"Failed to connect to RMBG client: {e}")
+    rmbg_client = None
+
+# 4. Load EnvMaps (Assuming assets folder exists)
+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.")
 
 # ==========================================
 # Helper Functions
@@ -199,48 +230,90 @@ def start_session(req: gr.Request):
     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):
+    try:
+        user_dir = os.path.join(TMP_DIR, str(req.session_hash))
         shutil.rmtree(user_dir)
+    except:
+        pass
 
 def remove_background(input: Image.Image) -> Image.Image:
-    with tempfile.NamedTemporaryFile(suffix='.png') as f:
+    """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:
         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)
+        f_path = f.name
+    
+    try:
+        output_path = rmbg_client.predict(handle_file(f_path), api_name="/image")[0][0]
+        output = Image.open(output_path)
         return output
+    finally:
+        if os.path.exists(f_path):
+            os.remove(f_path)
 
 def preprocess_image(input: Image.Image) -> Image.Image:
-    """Preprocess the input image: remove bg, crop, resize."""
+    """Preprocesses image (resizing, centering, BG removal) for TRELLIS."""
+    if input is None:
+        return None
+        
+    # Check Alpha
     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)
-    alpha = output_np[:, :, 3]
-    bbox = np.argwhere(alpha > 0.8 * 255)
-    if bbox.size == 0:
-        return output # Return original if empty
-    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))
+    # 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))
     return output
 
 def pack_state(latents: Tuple[SparseTensor, SparseTensor, int]) -> dict:
@@ -264,34 +337,30 @@ def get_seed(randomize_seed: bool, seed: int) -> int:
     return np.random.randint(0, MAX_SEED) if randomize_seed else seed
 
 # ==========================================
-# Inference Logic
+# Main Processing Functions
 # ==========================================
 
 @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.")
+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.")
     if not prompt.strip():
-        raise gr.Error("Please enter a prompt.")
-    
+        raise gr.Error("Prompt is empty.")
+
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    generator = torch.Generator(device).manual_seed(42) # Or random
+    generator = torch.Generator(device).manual_seed(42) # Fixed seed for consistency demo, or make parametric
     
-    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,
-        )
-        return result.images[0]
-    except Exception as e:
-        raise gr.Error(f"Z-Image Generation failed: {str(e)}")
+    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]
 
 @spaces.GPU(duration=120)
 def image_to_3d(
@@ -312,16 +381,24 @@ def image_to_3d(
     tex_slat_rescale_t: float,
     req: gr.Request,
     progress=gr.Progress(track_tqdm=True),
-) -> str:
+) -> Tuple[dict, str]:
     
+    if pipeline_trellis is None:
+        raise gr.Error("TRELLIS Model not loaded.")
+
     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.
+
     # --- Sampling ---
-    outputs, latents = trellis_pipeline.run(
+    outputs, latents = pipeline_trellis.run(
         image,
         seed=seed,
-        preprocess_image=False, # We pre-process in the upload handler or assume clean input
+        preprocess_image=False, # We assume input is already preprocessed
         sparse_structure_sampler_params={
             "steps": ss_sampling_steps,
             "guidance_strength": ss_guidance_strength,
@@ -340,45 +417,71 @@ def image_to_3d(
             "guidance_rescale": tex_slat_guidance_rescale,
             "rescale_t": tex_slat_rescale_t,
         },
-        pipeline_type={"512": "512", "1024": "1024_cascade", "1536": "1536_cascade"}[resolution],
+        pipeline_type={
+            "512": "512",
+            "1024": "1024_cascade",
+            "1536": "1536_cascade",
+        }[resolution],
         return_latent=True,
     )
+    
     mesh = outputs[0]
     mesh.simplify(16777216) 
     
     # Render Preview
-    images = render_utils.render_snapshot(mesh, resolution=1024, r=2, fov=36, nviews=STEPS, envmap=envmap)
+    images_rendered = 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[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_rendered[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):        
+    for idx, mode in enumerate(MODES): 
+        if mode['render_key'] not in images_rendered: continue
         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 modes.</p><p>● <b>View Angle</b> - Drag slider to rotate.</p></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>
-        <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)
@@ -389,12 +492,10 @@ 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 = trellis_pipeline.decode_latent(shape_slat, tex_slat, res)[0]
+    mesh = pipeline_trellis.decode_latent(shape_slat, tex_slat, res)[0]
     mesh.simplify(16777216)
     
     glb = o_voxel.postprocess.to_glb(
@@ -402,7 +503,7 @@ def extract_glb(
         faces=mesh.faces,
         attr_volume=mesh.attrs,
         coords=mesh.coords,
-        attr_layout=trellis_pipeline.pbr_attr_layout,
+        attr_layout=pipeline_trellis.pbr_attr_layout,
         grid_size=res,
         aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
         decimation_target=decimation_target,
@@ -412,6 +513,7 @@ 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)
@@ -421,113 +523,109 @@ def extract_glb(
     return glb_path, glb_path
 
 # ==========================================
-# Gradio UI Blocks
+# Gradio Interface
 # ==========================================
 
 if __name__ == "__main__":
     os.makedirs(TMP_DIR, exist_ok=True)
     
-    # Pre-process icon base64
+    # 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'] = "" # Fallback
+            MODES[i]['icon_base64'] = "" # Handle missing assets
 
-    with gr.Blocks(delete_cache=(600, 600)) as demo:
+    with gr.Blocks(css=CSS, head=HEAD, delete_cache=(600, 600)) as demo:
         gr.Markdown("""
         # TRELLIS.2-3D
-        **Unified Text-to-3D Workflow**
-        
-        1. **Text to Image**: Generate a base image using Z-Image-Turbo.
-        2. **Image to 3D**: Convert that image into a high-quality 3D asset using TRELLIS.2.
+        ### Text-to-Image (Z-Image-Turbo) + Image-to-3D (TRELLIS.2) Pipeline
         """)
         
         with gr.Row():
-            # --- Column 1: Inputs & Config ---
-            with gr.Column(scale=1, min_width=360):
-                
-                # --- Step 1: Text to Image ---
-                with gr.Group():
-                    gr.Markdown("### Step 1: Generate Image")
-                    txt_prompt = gr.Textbox(label="Prompt", placeholder="A sci-fi helmet, high quality, white background")
-                    btn_gen_img = gr.Button("Generate Image", variant="secondary")
+            # 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")
                 
-                # --- Step 2: Image to 3D Input ---
-                gr.Markdown("### Step 2: Configure & Convert")
-                image_prompt = gr.Image(label="Input Image (Generated or Uploaded)", format="png", image_mode="RGBA", type="pil", height=300)
+                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)
                 
-                with gr.Accordion("3D Generation Settings", open=True):
+                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)
-                    decimation_target = gr.Slider(100000, 500000, label="Decimation Target (For GLB)", value=300000, step=10000)
-                    texture_size = gr.Slider(1024, 4096, label="Texture Size (For GLB)", value=2048, step=1024)
-                
-                btn_gen_3d = gr.Button("Generate 3D", variant="primary")
-
-                with gr.Accordion(label="Advanced Sampling Settings", 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")
-
-            # --- Column 2: Outputs ---
-            with gr.Column(scale=10):
+                    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 Asset Preview", show_label=True, container=True)
-                        extract_btn = gr.Button("Extract GLB")
-                    with gr.Step("Extract", id=1):
-                        glb_output = gr.Model3D(label="Extracted GLB", height=724, show_label=True, display_mode="solid", clear_color=(0.25, 0.25, 0.25, 1.0))
-                        download_btn = gr.DownloadButton(label="Download GLB")
-                
-                gr.Markdown("Note: GLB extraction might take ~30s.")
-
-        # ==========================================
-        # Wiring Events
-        # ==========================================
+                        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
+        # ====================
         
-        # State to hold the latent 3D representation
-        output_buf = gr.State()
-
         demo.load(start_session)
         demo.unload(end_session)
 
-        # 1. Text to Image Event
-        btn_gen_img.click(
-            generate_txt2img,
-            inputs=[txt_prompt],
-            outputs=[image_prompt]
-        ).then(
-            preprocess_image, # Auto preprocess the generated image (rmbg)
-            inputs=[image_prompt],
-            outputs=[image_prompt]
+        # 1. Text to Image
+        gen_image_btn.click(
+            generate_text_to_image,
+            inputs=[text_prompt],
+            outputs=[image_input]
         )
 
-        # 2. Upload Image Event (Preprocess)
-        image_prompt.upload(
+        # 2. Image Preprocessing (Auto-trigger when image changes)
+        image_input.change(
             preprocess_image,
-            inputs=[image_prompt],
-            outputs=[image_prompt],
+            inputs=[image_input],
+            outputs=[image_input]
         )
 
-        # 3. Image to 3D Event
-        btn_gen_3d.click(
+        # 3. Image to 3D Generation
+        gen_3d_btn.click(
             get_seed,
             inputs=[randomize_seed, seed],
             outputs=[seed],
@@ -536,21 +634,24 @@ if __name__ == "__main__":
         ).then(
             image_to_3d,
             inputs=[
-                image_prompt, seed, resolution,
+                image_input, seed, resolution,
                 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,
+                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],
+            outputs=[output_state, preview_output],
         )
 
-        # 4. Extraction Event
+        # 4. Extract GLB
         extract_btn.click(
             lambda: gr.Walkthrough(selected=1), outputs=walkthrough
         ).then(
             extract_glb,
-            inputs=[output_buf, decimation_target, texture_size],
+            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)
 
-    demo.launch(css=css, head=head, mcp_server=True, ssr_mode=False, show_error=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860)