update app [.]

app.py

@@ -1,44 +1,45 @@
 import os
-import io
-import cv2
-import time
 import shutil
-import base64
+import cv2
 import torch
 import numpy as np
-import tempfile
-import gradio as gr
-import spaces
-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 gradio_client import Client, handle_file
+from pathlib import Path
 
-# --- 1. Environment Configuration ---
+# --- Environment Setup (Must be before other imports) ---
 os.environ["OPENCV_IO_ENABLE_OPENEXR"] = '1'
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 os.environ["ATTN_BACKEND"] = "flash_attn_3"
-# Note: Ensure autotune_cache.json exists or remove this line if not needed
+# 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')
 os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1'
 
-# --- 2. Imports from specific libraries ---
-# Ensure Z-Image-Turbo dependencies are met
+# --- Third Party Imports ---
+import gradio as gr
+from gradio_client import Client, handle_file
+import spaces
 from diffusers import ZImagePipeline
 
-# Ensure TRELLIS dependencies are met
+# --- TRELLIS Specific Imports ---
 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
 
-# --- 3. Constants & Settings ---
+# ==========================================
+# Global Configuration & Assets
+# ==========================================
+
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 
-# Rendering Modes configuration
+# TRELLIS Render Modes
 MODES = [
     {"name": "Normal", "icon": "assets/app/normal.png", "render_key": "normal"},
     {"name": "Clay render", "icon": "assets/app/clay.png", "render_key": "clay"},
@@ -51,8 +52,12 @@ STEPS = 8
 DEFAULT_MODE = 3
 DEFAULT_STEP = 3
 
-# --- 4. CSS & JavaScript (From TRELLIS) ---
+# ==========================================
+# CSS & JavaScript (For Custom Previewer)
+# ==========================================
+
 css = """
+/* Overwrite Gradio Default Style */
 .stepper-wrapper { padding: 0; }
 .stepper-container { padding: 0; align-items: center; }
 .step-button { flex-direction: row; }
@@ -61,23 +66,38 @@ css = """
 .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; bottom: 0 !important; right: 0 !important; transform: unset !important; }
-.previewer-container { position: relative; font-family: sans-serif; width: 100%; height: 722px; margin: 0 auto; padding: 20px; display: flex; flex-direction: column; align-items: center; justify-content: center; }
+.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;
+}
 .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-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 .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 */
 .gradio-container .padded:has(.previewer-container) { padding: 0 !important; }
 .gradio-container:has(.previewer-container) [data-testid="block-label"] { position: absolute; top: 0; left: 0; }
 """
@@ -100,6 +120,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');
@@ -111,58 +132,60 @@ head = """
 </script>
 """
 
-empty_html = """
+empty_html = f"""
 <div class="previewer-container">
-    <div style="text-align: center; color: var(--body-text-color); opacity: 0.5;">
-        <p>Generated 3D 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>
 """
 
-# --- 5. Global Model Initialization ---
-print("Initializing Models...")
+# ==========================================
+# Model Loading
+# ==========================================
+
+print("Initializing models...")
 
-# A. Load Z-Image-Turbo
+# 1. Z-Image-Turbo (Text to Image)
 print("Loading Z-Image-Turbo...")
-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")
-
-# B. Load TRELLIS.2
-print("Loading TRELLIS.2-4B...")
+try:
+    z_pipe = ZImagePipeline.from_pretrained(
+        "Tongyi-MAI/Z-Image-Turbo",
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=False,
+    )
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    z_pipe.to(device)
+    print("Z-Image-Turbo loaded.")
+except Exception as e:
+    print(f"Failed to load Z-Image-Turbo: {e}")
+    z_pipe = None
+
+# 2. 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  # We use external API/Client for robustness
+trellis_pipeline.rembg_model = None
 trellis_pipeline.low_vram = False
-if torch.cuda.is_available():
-    trellis_pipeline.cuda()
+trellis_pipeline.cuda()
 
-# C. Load RMBG Client (Background Removal)
+# 3. Background Remover
 rmbg_client = Client("briaai/BRIA-RMBG-2.0")
 
-# D. Load Environment Maps for Rendering
-# Helper to safely load local assets
-def load_env_map(name, filename):
-    path = os.path.join('assets/hdri', filename)
-    if os.path.exists(path):
-        return EnvMap(torch.tensor(
-            cv2.cvtColor(cv2.imread(path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB),
-            dtype=torch.float32, device='cuda'
-        ))
-    else:
-        print(f"Warning: HDRI {filename} not found. Rendering might fail or be black.")
-        return None # Should handle fallback
-
-envmap = {
-    'forest': load_env_map('forest', 'forest.exr'),
-    'sunset': load_env_map('sunset', 'sunset.exr'),
-    'courtyard': load_env_map('courtyard', 'courtyard.exr'),
-}
+# 4. HDRI Maps for TRELLIS
+envmap = {}
+# Try to load assets, handle gracefully if running in a basic environment
+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 HDRI assets. {e}")
 
-# --- 6. Helper Functions ---
+# ==========================================
+# Helper Functions
+# ==========================================
 
 def image_to_base64(image):
     buffered = io.BytesIO()
@@ -180,92 +203,89 @@ def end_session(req: gr.Request):
     if os.path.exists(user_dir):
         shutil.rmtree(user_dir)
 
-def pack_state(latents: Tuple[SparseTensor, SparseTensor, int]) -> dict:
-    shape_slat, tex_slat, res = latents
-    return {
-        'shape_slat_feats': shape_slat.feats.cpu().numpy(),
-        'tex_slat_feats': tex_slat.feats.cpu().numpy(),
-        '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(),
-        coords=torch.from_numpy(state['coords']).cuda(),
-    )
-    tex_slat = shape_slat.replace(torch.from_numpy(state['tex_slat_feats']).cuda())
-    return shape_slat, tex_slat, state['res']
-
-def get_seed(randomize_seed: bool, seed: int) -> int:
-    return np.random.randint(0, MAX_SEED) if randomize_seed else seed
-
-def remove_background(input_img: Image.Image) -> Image.Image:
-    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
-        input_img = input_img.convert('RGB')
-        input_img.save(f.name)
-        # Use gradio client to call remote or local rmbg
+def remove_background(input: Image.Image) -> Image.Image:
+    with tempfile.NamedTemporaryFile(suffix='.png') as f:
+        input = input.convert('RGB')
+        input.save(f.name)
         output = rmbg_client.predict(handle_file(f.name), api_name="/image")[0][0]
-        output_img = Image.open(output)
-        return output_img
+        output = Image.open(output)
+        return output
 
-def preprocess_image(input_img: Image.Image) -> Image.Image:
-    if input_img is None: return None
-    # Check alpha
+def preprocess_image(input: Image.Image) -> Image.Image:
+    """Preprocess the input image: remove bg, crop, resize."""
     has_alpha = False
-    if input_img.mode == 'RGBA':
-        alpha = np.array(input_img)[:, :, 3]
+    if input.mode == 'RGBA':
+        alpha = np.array(input)[:, :, 3]
         if not np.all(alpha == 255):
             has_alpha = True
-            
-    # Resize logic
-    max_size = max(input_img.size)
+    max_size = max(input.size)
     scale = min(1, 1024 / max_size)
     if scale < 1:
-        input_img = input_img.resize((int(input_img.width * scale), int(input_img.height * scale)), Image.Resampling.LANCZOS)
-    
+        input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS)
     if has_alpha:
-        output = input_img
+        output = input
     else:
-        output = remove_background(input_img)
-        
-    # Recenter and crop
+        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 # Fallback if empty
-    
+    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) # margin factor
+    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)
-    
-    # Premultiply alpha
     output = np.array(output).astype(np.float32) / 255
     output = output[:, :, :3] * output[:, :, 3:4]
     output = Image.fromarray((output * 255).astype(np.uint8))
     return output
 
-# --- 7. Core GPU Generators ---
+def pack_state(latents: Tuple[SparseTensor, SparseTensor, int]) -> dict:
+    shape_slat, tex_slat, res = latents
+    return {
+        'shape_slat_feats': shape_slat.feats.cpu().numpy(),
+        'tex_slat_feats': tex_slat.feats.cpu().numpy(),
+        '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(),
+        coords=torch.from_numpy(state['coords']).cuda(),
+    )
+    tex_slat = shape_slat.replace(torch.from_numpy(state['tex_slat_feats']).cuda())
+    return shape_slat, tex_slat, state['res']
+
+def get_seed(randomize_seed: bool, seed: int) -> int:
+    return np.random.randint(0, MAX_SEED) if randomize_seed else seed
+
+# ==========================================
+# Inference Logic
+# ==========================================
 
 @spaces.GPU()
-def generate_text_to_image(prompt):
-    """Generates image from text using Z-Image-Turbo."""
+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) # Fixed seed for T2I stability, or make random
+    generator = torch.Generator(device).manual_seed(42) # Or random
     
-    print(f"Generating image for: {prompt}")
+    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, # Z-Image-Turbo specific
+            num_inference_steps=9,
             guidance_scale=0.0,
             generator=generator,
         )
@@ -273,70 +293,65 @@ def generate_text_to_image(prompt):
     except Exception as e:
         raise gr.Error(f"Z-Image Generation failed: {str(e)}")
 
-@spaces.GPU(duration=180)
-def generate_3d_from_image(
+@spaces.GPU(duration=120)
+def image_to_3d(
     image: Image.Image,
     seed: int,
     resolution: str,
-    ss_params: list,
-    shape_params: list,
-    tex_params: list,
-):
-    """
-    Main pipeline: Preprocess Image -> Generate TRELLIS Latents -> Render Preview
-    """
-    if image is None:
-        raise gr.Error("No image provided for 3D generation.")
-
-    # 1. Preprocess (Remove BG if needed, center)
-    processed_image = preprocess_image(image)
+    ss_guidance_strength: float,
+    ss_guidance_rescale: float,
+    ss_sampling_steps: int,
+    ss_rescale_t: float,
+    shape_slat_guidance_strength: float,
+    shape_slat_guidance_rescale: float,
+    shape_slat_sampling_steps: int,
+    shape_slat_rescale_t: float,
+    tex_slat_guidance_strength: float,
+    tex_slat_guidance_rescale: float,
+    tex_slat_sampling_steps: int,
+    tex_slat_rescale_t: float,
+    req: gr.Request,
+    progress=gr.Progress(track_tqdm=True),
+) -> str:
     
-    # 2. Extract params
-    (ss_steps, ss_guidance, ss_rescale, ss_t) = ss_params
-    (sh_steps, sh_guidance, sh_rescale, sh_t) = shape_params
-    (tex_steps, tex_guidance, tex_rescale, tex_t) = tex_params
+    if image is None:
+        raise gr.Error("Input image is missing.")
 
-    # 3. Run TRELLIS Pipeline
+    # --- Sampling ---
     outputs, latents = trellis_pipeline.run(
-        processed_image,
+        image,
         seed=seed,
-        preprocess_image=False, # We did it manually
+        preprocess_image=False, # We pre-process in the upload handler or assume clean input
         sparse_structure_sampler_params={
-            "steps": int(ss_steps),
-            "guidance_strength": ss_guidance,
-            "guidance_rescale": ss_rescale,
-            "rescale_t": ss_t,
+            "steps": ss_sampling_steps,
+            "guidance_strength": ss_guidance_strength,
+            "guidance_rescale": ss_guidance_rescale,
+            "rescale_t": ss_rescale_t,
         },
         shape_slat_sampler_params={
-            "steps": int(sh_steps),
-            "guidance_strength": sh_guidance,
-            "guidance_rescale": sh_rescale,
-            "rescale_t": sh_t,
+            "steps": shape_slat_sampling_steps,
+            "guidance_strength": shape_slat_guidance_strength,
+            "guidance_rescale": shape_slat_guidance_rescale,
+            "rescale_t": shape_slat_rescale_t,
         },
         tex_slat_sampler_params={
-            "steps": int(tex_steps),
-            "guidance_strength": tex_guidance,
-            "guidance_rescale": tex_rescale,
-            "rescale_t": tex_t,
+            "steps": tex_slat_sampling_steps,
+            "guidance_strength": tex_slat_guidance_strength,
+            "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,
     )
-
-    # 4. Render Preview
     mesh = outputs[0]
-    mesh.simplify(16777216)
+    mesh.simplify(16777216) 
     
-    # Use environment map for rendering if available
+    # Render Preview
     images = render_utils.render_snapshot(mesh, resolution=1024, r=2, fov=36, nviews=STEPS, envmap=envmap)
     state = pack_state(latents)
     torch.cuda.empty_cache()
-
-    # 5. Build HTML
+    
+    # --- HTML Construction ---
     images_html = ""
     for m_idx, mode in enumerate(MODES):
         for s_idx in range(STEPS):
@@ -344,18 +359,18 @@ def generate_3d_from_image(
             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">'
-            
+            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):        
         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> - Switch render modes.</p><p>● <b>View Angle</b> - Drag slider.</p></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>
         <div class="display-row">{images_html}</div>
         <div class="mode-row" id="btn-group">{btns_html}</div>
@@ -364,14 +379,21 @@ def generate_3d_from_image(
         </div>
     </div>
     """
-    
-    return state, full_html, processed_image
+    return state, full_html
 
 @spaces.GPU(duration=120)
-def extract_glb(state: dict, decimation_target: int, texture_size: int, req: gr.Request):
+def extract_glb(
+    state: dict,
+    decimation_target: int,
+    texture_size: int,
+    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.simplify(16777216)
     
@@ -390,150 +412,145 @@ def extract_glb(state: dict, decimation_target: int, texture_size: int, req: gr.
         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)
     glb_path = os.path.join(user_dir, f'sample_{timestamp}.glb')
     glb.export(glb_path, extension_webp=True)
     torch.cuda.empty_cache()
-    
     return glb_path, glb_path
 
-# --- 8. UI Construction ---
+# ==========================================
+# Gradio UI Blocks
+# ==========================================
 
-# Load icons for UI (Assumes assets folder exists)
-for i in range(len(MODES)):
-    icon_path = MODES[i]['icon']
-    if os.path.exists(icon_path):
-        MODES[i]['icon_base64'] = image_to_base64(Image.open(icon_path))
-    else:
-        # Fallback empty image if asset missing
-        MODES[i]['icon_base64'] = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
-
-with gr.Blocks(css=css, title="TRELLIS.2-3D Turbo", delete_cache=(600, 600)) as demo:
+if __name__ == "__main__":
+    os.makedirs(TMP_DIR, exist_ok=True)
     
-    # Session Management
-    demo.load(start_session)
-    demo.unload(end_session)
-
-    gr.Markdown("""
-    # ⚡ Text to 3D Turbo
-    **Z-Image-Turbo** (Text-to-Image) + **TRELLIS.2** (Image-to-3D)
-    """)
-
-    with gr.Row():
-        # LEFT COLUMN: INPUTS
-        with gr.Column(scale=1, min_width=350):
-            text_prompt = gr.Textbox(label="Text Prompt", placeholder="A stylized cute dragon, high quality, 3d render...")
-            generate_txt2img_btn = gr.Button("1. Generate Image", variant="primary")
-            
-            # Intermediate Image
-            generated_image = gr.Image(label="Generated Reference Image", type="pil", height=300)
-            
-            gr.Markdown("---")
-            generate_3d_btn = gr.Button("2. Generate 3D Model", variant="secondary")
-            
-            with gr.Accordion("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)
+    # Pre-process icon base64
+    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
+
+    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.
+        """)
+        
+        with gr.Row():
+            # --- Column 1: Inputs & Config ---
+            with gr.Column(scale=1, min_width=360):
                 
-                # Grouped params for passing to function
-                # SS
-                ss_steps = gr.Slider(1, 50, value=12, label="Structure Steps")
-                ss_str = gr.Slider(1.0, 10.0, value=7.5, label="Structure Guidance")
-                ss_rescale = gr.Slider(0.0, 1.0, value=0.7, label="Structure Rescale")
-                ss_t = gr.Slider(1.0, 6.0, value=5.0, label="Structure T")
+                # --- 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")
                 
-                # Shape
-                sh_steps = gr.Slider(1, 50, value=12, label="Shape Steps")
-                sh_str = gr.Slider(1.0, 10.0, value=7.5, label="Shape Guidance")
-                sh_rescale = gr.Slider(0.0, 1.0, value=0.5, label="Shape Rescale")
-                sh_t = gr.Slider(1.0, 6.0, value=3.0, label="Shape T")
+                # --- 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)
                 
-                # Texture
-                tex_steps = gr.Slider(1, 50, value=12, label="Texture Steps")
-                tex_str = gr.Slider(1.0, 10.0, value=1.0, label="Texture Guidance")
-                tex_rescale = gr.Slider(0.0, 1.0, value=0.0, label="Texture Rescale")
-                tex_t = gr.Slider(1.0, 6.0, value=3.0, label="Texture T")
-
-        # RIGHT COLUMN: OUTPUTS
-        with gr.Column(scale=2):
-            with gr.Tabs():
-                with gr.Tab("3D Preview"):
-                    preview_html = gr.HTML(empty_html, label="3D Preview")
-                    # Hidden state to store latent representation
-                    trellis_state = gr.State()
+                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)
+                    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)
                 
-                with gr.Tab("Export GLB"):
-                    gr.Markdown("Extract GLB from the generated 3D data.")
-                    decimation = gr.Slider(100000, 500000, label="Decimation Target", value=300000, step=10000)
-                    tex_size = gr.Slider(1024, 4096, label="Texture Size", value=2048, step=1024)
-                    extract_btn = gr.Button("Extract & Download GLB")
-                    glb_model = gr.Model3D(label="Extracted GLB", display_mode="solid", clear_color=(0.2, 0.2, 0.2, 1.0))
-                    dl_button = gr.DownloadButton("Download .glb")
-
-    # --- Wiring ---
-    
-    # 1. Text to Image
-    generate_txt2img_btn.click(
-        generate_text_to_image,
-        inputs=[text_prompt],
-        outputs=[generated_image]
-    )
+                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.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.")
 
-    # 2. Image to 3D
-    generate_3d_btn.click(
-        get_seed, inputs=[randomize_seed, seed], outputs=[seed]
-    ).then(
-        generate_3d_from_image,
-        inputs=[
-            generated_image, seed, resolution,
-            # Pack params into lists to keep signature clean
-            gr.State([12, 7.5, 0.7, 5.0]), # Default SS placeholder, normally would bind sliders directly
-            gr.State([12, 7.5, 0.5, 3.0]), # Default Shape placeholder
-            gr.State([12, 1.0, 0.0, 3.0])  # Default Tex placeholder
-        ], 
-        # Actually bind the real sliders
-        inputs_kwargs={
-            "ss_params": [ss_steps, ss_str, ss_rescale, ss_t],
-            "shape_params": [sh_steps, sh_str, sh_rescale, sh_t],
-            "tex_params": [tex_steps, tex_str, tex_rescale, tex_t]
-        },
-        outputs=[trellis_state, preview_html, generated_image] # Update image with processed version
-    )
+        # ==========================================
+        # Wiring Events
+        # ==========================================
+        
+        # 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]
+        )
 
-    # Fix: Gradio .then chaining with lists of inputs needs careful handling. 
-    # Let's redefine the click to map inputs directly without the inputs_kwargs workaround for clarity.
-    generate_3d_btn.click(
-        get_seed, inputs=[randomize_seed, seed], outputs=[seed]
-    ).then(
-        lambda *args: generate_3d_from_image(
-            args[0], args[1], args[2],
-            [args[3], args[4], args[5], args[6]], # SS params
-            [args[7], args[8], args[9], args[10]], # Shape params
-            [args[11], args[12], args[13], args[14]] # Tex params
-        ),
-        inputs=[
-            generated_image, seed, resolution,
-            ss_steps, ss_str, ss_rescale, ss_t,
-            sh_steps, sh_str, sh_rescale, sh_t,
-            tex_steps, tex_str, tex_rescale, tex_t
-        ],
-        outputs=[trellis_state, preview_html, generated_image]
-    )
+        # 2. Upload Image Event (Preprocess)
+        image_prompt.upload(
+            preprocess_image,
+            inputs=[image_prompt],
+            outputs=[image_prompt],
+        )
 
-    # 3. Extract GLB
-    extract_btn.click(
-        extract_glb,
-        inputs=[trellis_state, decimation, tex_size],
-        outputs=[glb_model, dl_button]
-    )
+        # 3. Image to 3D Event
+        btn_gen_3d.click(
+            get_seed,
+            inputs=[randomize_seed, seed],
+            outputs=[seed],
+        ).then(
+            lambda: gr.Walkthrough(selected=0), outputs=walkthrough
+        ).then(
+            image_to_3d,
+            inputs=[
+                image_prompt, 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,
+            ],
+            outputs=[output_buf, preview_output],
+        )
 
-if __name__ == "__main__":
-    # Create temp directory
-    os.makedirs(TMP_DIR, exist_ok=True)
-    
-    # Launch with custom scripts
-    demo.launch(head=head)
+        # 4. Extraction Event
+        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],
+        )
+
+    demo.launch()