Update app.py

app.py

@@ -1,11 +1,8 @@
 import os
 import shutil
-import cv2
 import torch
 import numpy as np
 from PIL import Image
-import base64
-import io
 import tempfile
 from typing import *
 from datetime import datetime
@@ -26,121 +23,18 @@ import spaces
 from diffusers import ZImagePipeline
 
 # --- 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
 
 # ==========================================
-# Global Configuration & Assets
+# Global Configuration
 # ==========================================
 
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 
-# 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"},
-    {"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 & 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; }
-.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;
-}
-.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; }
-"""
-
-head = """
-<script>
-    function refreshView(mode, step) {
-        const allImgs = document.querySelectorAll('.previewer-main-image');
-        for (let i = 0; i < allImgs.length; i++) {
-            const img = allImgs[i];
-            if (img.classList.contains('visible')) {
-                const id = img.id;
-                const [_, m, s] = id.split('-');
-                if (mode === -1) mode = parseInt(m.slice(1));
-                if (step === -1) step = parseInt(s.slice(1));
-                break;
-            }
-        }
-        allImgs.forEach(img => img.classList.remove('visible'));
-        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');
-            else btn.classList.remove('active');
-        });
-    }
-    function selectMode(mode) { refreshView(mode, -1); }
-    function onSliderChange(val) { refreshView(-1, parseInt(val)); }
-</script>
-"""
-
-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>
-"""
-
-# ==========================================
-# Model & Asset Loading
+# Model Loading
 # ==========================================
 
 print("Initializing models...")
@@ -162,49 +56,23 @@ except Exception as e:
 
 # 2. TRELLIS.2 (Image to 3D)
 print("Loading TRELLIS.2...")
-trellis_pipeline = Trellis2ImageTo3DPipeline.from_pretrained('microsoft/TRELLIS.2-4B')
-trellis_pipeline.rembg_model = None
-trellis_pipeline.low_vram = False
-trellis_pipeline.cuda()
+try:
+    trellis_pipeline = Trellis2ImageTo3DPipeline.from_pretrained('microsoft/TRELLIS.2-4B')
+    trellis_pipeline.rembg_model = None
+    trellis_pipeline.low_vram = False
+    trellis_pipeline.cuda()
+    print("TRELLIS.2 loaded.")
+except Exception as e:
+    print(f"Failed to load TRELLIS.2: {e}")
+    trellis_pipeline = None
 
 # 3. Background Remover
 rmbg_client = Client("briaai/BRIA-RMBG-2.0")
 
-# 4. HDRI Maps (FIXED: Robust Fallback)
-def load_envmap_safe(path):
-    """Attempt to load an HDRI map, return a dummy white map if it fails."""
-    try:
-        if not os.path.exists(path):
-            raise FileNotFoundError(f"{path} not found")
-        im = cv2.imread(path, cv2.IMREAD_UNCHANGED)
-        if im is None:
-            raise ValueError(f"OpenCV returned None for {path}")
-        rgb = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
-        return EnvMap(torch.tensor(rgb, dtype=torch.float32, device='cuda'))
-    except Exception as e:
-        print(f"Warning: Could not load HDRI from {path} ({e}). Using synthetic fallback.")
-        # Create a simple white environment map (16x32 resolution)
-        synthetic_env = torch.ones((16, 32, 3), dtype=torch.float32, device='cuda')
-        return EnvMap(synthetic_env)
-
-# Always populate these keys, even if files are missing
-envmap = {
-    'forest': load_envmap_safe('assets/hdri/forest.exr'),
-    'sunset': load_envmap_safe('assets/hdri/sunset.exr'),
-    'courtyard': load_envmap_safe('assets/hdri/courtyard.exr'),
-}
-
 # ==========================================
 # Helper Functions
 # ==========================================
 
-def image_to_base64(image):
-    buffered = io.BytesIO()
-    image = image.convert("RGB")
-    image.save(buffered, format="jpeg", quality=85)
-    img_str = base64.b64encode(buffered.getvalue()).decode()
-    return f"data:image/jpeg;base64,{img_str}"
-
 def start_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     os.makedirs(user_dir, exist_ok=True)
@@ -257,23 +125,6 @@ def preprocess_image(input: Image.Image) -> Image.Image:
     output = Image.fromarray((output * 255).astype(np.uint8))
     return output
 
-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
 
@@ -290,7 +141,7 @@ def generate_txt2img(prompt, progress=gr.Progress(track_tqdm=True)):
         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 consistency demo
+    generator = torch.Generator(device).manual_seed(42)
     
     progress(0.1, desc="Generating Text-to-Image...")
     try:
@@ -307,131 +158,85 @@ def generate_txt2img(prompt, progress=gr.Progress(track_tqdm=True)):
     except Exception as e:
         raise gr.Error(f"Z-Image Generation failed: {str(e)}")
 
-@spaces.GPU(duration=120)
-def image_to_3d(
+@spaces.GPU(duration=300)
+def generate_3d(
     image: Image.Image,
     seed: int,
     resolution: str,
+    decimation_target: int,
+    texture_size: int,
     ss_guidance_strength: float,
     ss_guidance_rescale: float,
     ss_sampling_steps: int,
     ss_rescale_t: float,
-    shape_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,
+    shape_guidance: float,
+    shape_rescale: float,
+    shape_steps: int,
+    shape_rescale_t: float,
+    tex_guidance: float,
+    tex_rescale: float,
+    tex_steps: int,
+    tex_rescale_t: float,
     req: gr.Request,
     progress=gr.Progress(track_tqdm=True),
-) -> str:
+) -> Tuple[str, str]:
     
     if image is None:
-        raise gr.Error("Input image is missing.")
-
-    # --- Sampling ---
-    outputs, latents = trellis_pipeline.run(
-        image,
-        seed=seed,
-        preprocess_image=False, # We pre-process in the upload handler or assume clean input
-        sparse_structure_sampler_params={
-            "steps": ss_sampling_steps,
-            "guidance_strength": ss_guidance_strength,
-            "guidance_rescale": ss_guidance_rescale,
-            "rescale_t": ss_rescale_t,
-        },
-        shape_slat_sampler_params={
-            "steps": shape_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": 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],
-        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)
-    state = pack_state(latents)
-    torch.cuda.empty_cache()
+        raise gr.Error("Please provide an input image.")
     
-    # --- HTML Construction ---
-    images_html = ""
-    for m_idx, mode in enumerate(MODES):
-        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 ""
-            
-            render_key = mode['render_key']
-            if render_key in images:
-                img_data = images[render_key][s_idx]
-                img_base64 = image_to_base64(Image.fromarray(img_data))
-                images_html += f"""<img id="{unique_id}" class="previewer-main-image {vis_class}" src="{img_base64}" loading="eager">"""
-            else:
-                images_html += f"""<div id="{unique_id}" class="previewer-main-image {vis_class}" style="width:100%;height:100%;background:#333;color:white;display:flex;align-items:center;justify-content:center;">Render Error</div>"""
-    
-    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']}">"""
-    
-    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>
-        <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
-
-# Increased timeout to 300s (5mins) for robust extraction
-@spaces.GPU(duration=300)
-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.")
-    
-    # Clear cache before heavy operation
-    torch.cuda.empty_cache()
-        
+    if trellis_pipeline is None:
+        raise gr.Error("TRELLIS model is not loaded.")
+
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    shape_slat, tex_slat, res = unpack_state(state)
+    os.makedirs(user_dir, exist_ok=True)
     
+    # 1. Run Pipeline
+    progress(0.1, desc="Generating 3D Geometry...")
     try:
-        # Decode Latents
-        mesh = trellis_pipeline.decode_latent(shape_slat, tex_slat, res)[0]
-        mesh.simplify(16777216)
+        outputs, latents = trellis_pipeline.run(
+            image,
+            seed=seed,
+            preprocess_image=False, # Assumed already preprocessed by UI handler
+            sparse_structure_sampler_params={
+                "steps": ss_sampling_steps,
+                "guidance_strength": ss_guidance_strength,
+                "guidance_rescale": ss_guidance_rescale,
+                "rescale_t": ss_rescale_t,
+            },
+            shape_slat_sampler_params={
+                "steps": shape_steps,
+                "guidance_strength": shape_guidance,
+                "guidance_rescale": shape_rescale,
+                "rescale_t": shape_rescale_t,
+            },
+            tex_slat_sampler_params={
+                "steps": tex_steps,
+                "guidance_strength": tex_guidance,
+                "guidance_rescale": tex_rescale,
+                "rescale_t": tex_rescale_t,
+            },
+            pipeline_type={"512": "512", "1024": "1024_cascade", "1536": "1536_cascade"}[resolution],
+            return_latent=True,
+        )
+        
+        # 2. Process Mesh
+        progress(0.7, desc="Processing Mesh...")
+        mesh = outputs[0]
+        mesh.simplify(16777216) # Simplify for processing limits
+        
+        # 3. Export to GLB
+        progress(0.9, desc="Baking Texture & Exporting GLB...")
+        
+        # Note: We use the latent grid resolution from the pipeline output
+        grid_size = latents[2] 
         
-        # Convert to GLB
         glb = o_voxel.postprocess.to_glb(
             vertices=mesh.vertices,
             faces=mesh.faces,
             attr_volume=mesh.attrs,
             coords=mesh.coords,
             attr_layout=trellis_pipeline.pbr_attr_layout,
-            grid_size=res,
+            grid_size=grid_size,
             aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
             decimation_target=decimation_target,
             texture_size=texture_size,
@@ -441,19 +246,19 @@ def extract_glb(
             use_tqdm=True,
         )
         
-        # Export
         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')
+        timestamp = now.strftime("%Y-%m-%dT%H%M%S")
+        glb_path = os.path.join(user_dir, f'trellis_output_{timestamp}.glb')
         glb.export(glb_path, extension_webp=True)
         
-    except Exception as e:
+        # Clean up
         torch.cuda.empty_cache()
-        raise gr.Error(f"Extraction failed: {str(e)}")
+        return glb_path, glb_path
         
-    torch.cuda.empty_cache()
-    return glb_path, glb_path
+    except Exception as e:
+        torch.cuda.empty_cache()
+        raise gr.Error(f"Generation failed: {str(e)}")
+
 
 # ==========================================
 # Gradio UI Blocks
@@ -461,50 +266,46 @@ def extract_glb(
 
 if __name__ == "__main__":
     os.makedirs(TMP_DIR, exist_ok=True)
-    
-    # Pre-process icon base64
-    for i in range(len(MODES)):
-        icon_path = MODES[i]['icon']
-        if os.path.exists(icon_path):
-            icon = Image.open(icon_path)
-            MODES[i]['icon_base64'] = image_to_base64(icon)
-        else:
-            MODES[i]['icon_base64'] = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
 
     with gr.Blocks(delete_cache=(600, 600)) as demo:
         gr.Markdown("""
-        # TRELLIS.2-3D
-        **Unified Text-to-3D Workflow**
+        # TRELLIS.2-3D (Direct Export)
         
-        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.
+        **Workflow:**
+        1. **Generate Image** (Text-to-Image) or **Upload Image**.
+        2. Click **Generate 3D** to create the GLB asset directly.
         """)
         
         with gr.Row():
-            # --- Column 1: Inputs & Config ---
+            # --- Column 1: Inputs ---
             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")
+                # Input Source
+                with gr.Tabs():
+                    with gr.Tab("Text to Image"):
+                        txt_prompt = gr.Textbox(label="Prompt", placeholder="A treasure chest, isometric style, white background", lines=3)
+                        btn_gen_img = gr.Button("Generate Image", variant="secondary")
+                    with gr.Tab("Upload"):
+                        gr.Markdown("Upload an image directly if you have one.")
                 
-                # --- 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)
+                # Image Display
+                image_prompt = gr.Image(label="Input Image", 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")
+                # 3D Settings
+                gr.Markdown("### 3D Settings")
+                with gr.Group():
+                    resolution = gr.Radio(["512", "1024", "1536"], label="Generation Resolution", value="1024")
                     seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
                     randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-                    # LOWERED DEFAULTS TO PREVENT BROWSER FREEZE
-                    decimation_target = gr.Slider(50000, 500000, label="Decimation Target (Faces)", value=100000, step=10000)
+                    
+                    gr.Markdown(" **Export Settings**")
+                    decimation_target = gr.Slider(50000, 500000, label="Target Faces", value=150000, step=10000)
                     texture_size = gr.Slider(512, 4096, label="Texture Size", value=1024, step=512)
                 
-                btn_gen_3d = gr.Button("Generate 3D", variant="primary")
+                btn_gen_3d = gr.Button("Generate 3D Asset", variant="primary", scale=2)
 
-                with gr.Accordion(label="Advanced Sampling Settings", open=False):                
+                # Advanced Settings
+                with gr.Accordion(label="Advanced Sampler 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")
@@ -524,28 +325,25 @@ if __name__ == "__main__":
                     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):
-                        # Lower default height to prevent full screen takeover
-                        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")
-                
-                gr.Markdown("Note: Extraction may take 1-2 minutes. If browser slows down, reduce 'Decimation Target'.")
+            with gr.Column(scale=2):
+                gr.Markdown("### 3D Output")
+                glb_output = gr.Model3D(
+                    label="Generated GLB", 
+                    display_mode="solid", 
+                    clear_color=(0.2, 0.2, 0.2, 1.0),
+                    height=600,
+                    interactive=True
+                )
+                download_btn = gr.DownloadButton(label="Download GLB File")
 
         # ==========================================
         # Wiring Events
         # ==========================================
-        
-        output_buf = gr.State()
 
         demo.load(start_session)
         demo.unload(end_session)
 
-        # 1. Text to Image Event
+        # 1. Text to Image
         btn_gen_img.click(
             generate_txt2img,
             inputs=[txt_prompt],
@@ -556,38 +354,28 @@ if __name__ == "__main__":
             outputs=[image_prompt]
         )
 
-        # 2. Upload Image Event (Preprocess)
+        # 2. Auto-preprocess uploaded images
         image_prompt.upload(
             preprocess_image,
             inputs=[image_prompt],
             outputs=[image_prompt],
         )
 
-        # 3. Image to 3D Event
+        # 3. Generate 3D
         btn_gen_3d.click(
             get_seed,
             inputs=[randomize_seed, seed],
             outputs=[seed],
         ).then(
-            lambda: gr.Walkthrough(selected=0), outputs=walkthrough
-        ).then(
-            image_to_3d,
+            generate_3d,
             inputs=[
-                image_prompt, seed, resolution,
+                image_prompt, seed, resolution, 
+                decimation_target, texture_size,
                 ss_guidance_strength, ss_guidance_rescale, ss_sampling_steps, ss_rescale_t,
                 shape_guidance, shape_rescale, shape_steps, shape_rescale_t,
                 tex_guidance, tex_rescale, tex_steps, tex_rescale_t,
             ],
-            outputs=[output_buf, preview_output],
-        )
-
-        # 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(css=css, head=head)
+    demo.launch()