app.py

import os
import shutil
import torch
import numpy as np
from PIL import Image
import tempfile
import uuid
import io
import json
import random
import string
import zipfile
from typing import *
from datetime import datetime
from pathlib import Path

# --- Environment Setup ---
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = '1'
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
os.environ["ATTN_BACKEND"] = "flash_attn_2"
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'

# --- Video Rendering: Try pyrender with different backends ---
HAS_PYRENDER = False
PYRENDER_BACKEND = None

for _backend in ['egl', 'osmesa']:
    try:
        os.environ['PYOPENGL_PLATFORM'] = _backend
        import pyrender
        HAS_PYRENDER = True
        PYRENDER_BACKEND = _backend
        print(f"pyrender loaded with {_backend} backend for video rendering.")
        break
    except Exception as e:
        print(f"pyrender import failed with {_backend}: {e}")

if not HAS_PYRENDER:
    print("pyrender not available. Will use matplotlib for video rendering.")

import trimesh
import cv2

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

# --- Theme and UI ---
from typing import Iterable
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

import gradio as gr
from gradio_client import Client, handle_file
import spaces
from diffusers import DiffusionPipeline
from trellis2.pipelines import Trellis2ImageTo3DPipeline
import o_voxel

# --- Rerun ---
import rerun as rr
try:
    import rerun.blueprint as rrb
except ImportError:
    rrb = None

from gradio_rerun import Rerun

MAX_SEED = np.iinfo(np.int32).max
TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')

# --- Theme Configuration ---
colors.orange_red = colors.Color(
    name="orange_red",
    c50="#FFF0E5",
    c100="#FFE0CC",
    c200="#FFC299",
    c300="#FFA366",
    c400="#FF8533",
    c500="#FF4500",
    c600="#E63E00",
    c700="#CC3700",
    c800="#B33000",
    c900="#992900",
    c950="#802200",
)

class OrangeRedTheme(Soft):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.gray,
        secondary_hue: colors.Color | str = colors.orange_red,
        neutral_hue: colors.Color | str = colors.slate,
        text_size: sizes.Size | str = sizes.text_lg,
        font: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
        ),
        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue,
            secondary_hue=secondary_hue,
            neutral_hue=neutral_hue,
            text_size=text_size,
            font=font,
            font_mono=font_mono,
        )
        super().set(
            background_fill_primary="*primary_50",
            background_fill_primary_dark="*primary_900",
            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
            button_primary_text_color="white",
            button_primary_text_color_hover="white",
            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
            button_secondary_text_color="black",
            button_secondary_text_color_hover="white",
            button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)",
            button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)",
            button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)",
            button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)",
            slider_color="*secondary_500",
            slider_color_dark="*secondary_600",
            block_title_text_weight="600",
            block_border_width="3px",
            block_shadow="*shadow_drop_lg",
            button_primary_shadow="*shadow_drop_lg",
            button_large_padding="11px",
            color_accent_soft="*primary_100",
            block_label_background_fill="*primary_200",
        )

orange_red_theme = OrangeRedTheme()

# --- Model Loading ---
print("Initializing models...")

print("Loading Z-Image-Turbo...")
try:
    z_pipe = DiffusionPipeline.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

print("Loading TRELLIS.2...")
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

rmbg_client = Client("briaai/BRIA-RMBG-2.0")

# --- Session Management ---
def start_session(req: gr.Request):
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)

def end_session(req: gr.Request):
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    if os.path.exists(user_dir):
        shutil.rmtree(user_dir)

# --- Background Removal & Preprocessing ---
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 = Image.open(output)
        return output

def preprocess_image(input: Image.Image) -> Image.Image:
    """Preprocess the input image: remove bg, crop, resize."""
    if input is None:
        return None

    has_alpha = False
    if input.mode == 'RGBA':
        alpha = np.array(input)[:, :, 3]
        if not np.all(alpha == 255):
            has_alpha = True
    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)
    if has_alpha:
        output = input
    else:
        output = remove_background(input)

    output_np = np.array(output)
    alpha = output_np[:, :, 3]
    bbox = np.argwhere(alpha > 0.8 * 255)
    if bbox.size == 0:
        return output
    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))
    return output

def get_seed(randomize_seed: bool, seed: int) -> int:
    return np.random.randint(0, MAX_SEED) if randomize_seed else seed

# --- Text-to-Image ---
@spaces.GPU(size="xlarge", duration=120)
def generate_txt2img(prompt, progress=gr.Progress(track_tqdm=True)):
    if z_pipe is None:
        raise gr.Error("Z-Image-Turbo model failed to load.")
    if not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    generator = torch.Generator(device).manual_seed(42)

    progress(0.1, desc="Generating Text-to-Image...")
    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)}")

# --- 3D Generation ---
@spaces.GPU(size="xlarge", duration=120)
def generate_3d(
    image: Image.Image,
    seed: int,
    resolution: str,
    decimation_target: int,
    texture_size: int,
    ss_guidance_strength: float,
    ss_guidance_rescale: float,
    ss_sampling_steps: int,
    ss_rescale_t: float,
    shape_guidance: float,
    shape_rescale: float,
    shape_steps: int,
    shape_rescale_t: float,
    tex_guidance: float,
    tex_rescale: float,
    tex_steps: int,
    tex_rescale_t: float,
    req: gr.Request,
    progress=gr.Progress(track_tqdm=True),
) -> Tuple[str, str, str, str]:
    """Returns: rrd_path, glb_path, glb_path_for_state, image_path_for_state"""

    if image is None:
        raise gr.Error("Please provide an input image.")

    if trellis_pipeline is None:
        raise gr.Error("TRELLIS model is not loaded.")

    session_hash = req.session_hash if req is not None else "default_session"
    user_dir = os.path.join(TMP_DIR, session_hash)
    os.makedirs(user_dir, exist_ok=True)

    # Save input image for dataset export
    now = datetime.now()
    timestamp = now.strftime("%Y-%m-%dT%H%M%S")
    image_path = os.path.join(user_dir, f'input_image_{timestamp}.png')
    image.save(image_path)

    progress(0.1, desc="Generating 3D Geometry...")
    try:
        outputs, latents = trellis_pipeline.run(
            image,
            seed=seed,
            preprocess_image=False,
            sparse_structure_sampler_params={
                "steps": ss_sampling_steps,
                "guidance_strength": ss_guidance_strength,
                "guidance_rescale": ss_guidance_rescale,
                "rescale_t": ss_rescale_t,
            },
            shape_slat_sampler_params={
                "steps": shape_steps,
                "guidance_strength": shape_guidance,
                "guidance_rescale": shape_rescale,
                "rescale_t": shape_rescale_t,
            },
            tex_slat_sampler_params={
                "steps": tex_steps,
                "guidance_strength": tex_guidance,
                "guidance_rescale": tex_rescale,
                "rescale_t": tex_rescale_t,
            },
            pipeline_type={"512": "512", "1024": "1024_cascade", "1536": "1536_cascade"}[resolution],
            return_latent=True,
        )

        # Process Mesh
        progress(0.7, desc="Processing Mesh...")
        mesh = outputs[0]
        mesh.simplify(1000000)

        # Export to GLB
        progress(0.9, desc="Baking Texture & Exporting GLB...")

        grid_size = latents[2]

        try:
            glb = o_voxel.postprocess.to_glb(
                vertices=mesh.vertices,
                faces=mesh.faces,
                attr_volume=mesh.attrs,
                coords=mesh.coords,
                attr_layout=trellis_pipeline.pbr_attr_layout,
                grid_size=grid_size,
                aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
                decimation_target=decimation_target,
                texture_size=texture_size,
                remesh=True,
                remesh_band=1,
                remesh_project=0,
                use_tqdm=True,
            )
        except RuntimeError as e:
            print(f"Warning: Post-processing failed with remesh=True. Error: {e}")
            print("Retrying with remesh=False (Standard mesh generation)...")
            glb = o_voxel.postprocess.to_glb(
                vertices=mesh.vertices,
                faces=mesh.faces,
                attr_volume=mesh.attrs,
                coords=mesh.coords,
                attr_layout=trellis_pipeline.pbr_attr_layout,
                grid_size=grid_size,
                aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
                decimation_target=decimation_target,
                texture_size=texture_size,
                remesh=False,
                remesh_band=1,
                remesh_project=0,
                use_tqdm=True,
            )

        glb_path = os.path.join(user_dir, f'trellis_output_{timestamp}.glb')
        glb.export(glb_path, extension_webp=False)

        # --- Rerun Visualization ---
        progress(0.95, desc="Creating Rerun Visualization...")

        run_id = str(uuid.uuid4())

        rec = None
        if hasattr(rr, "new_recording"):
            rec = rr.new_recording(application_id="TRELLIS-3D-Viewer", recording_id=run_id)
        elif hasattr(rr, "RecordingStream"):
            rec = rr.RecordingStream(application_id="TRELLIS-3D-Viewer", recording_id=run_id)
        else:
            rr.init("TRELLIS-3D-Viewer", recording_id=run_id, spawn=False)
            rec = rr

        rec.log("world", rr.Clear(recursive=True), static=True)
        rec.log("world", rr.ViewCoordinates.RIGHT_HAND_Y_UP, static=True)

        try:
            rec.log("world/axes/x", rr.Arrows3D(vectors=[[0.5, 0, 0]], colors=[[255, 0, 0]]), static=True)
            rec.log("world/axes/y", rr.Arrows3D(vectors=[[0, 0.5, 0]], colors=[[0, 255, 0]]), static=True)
            rec.log("world/axes/z", rr.Arrows3D(vectors=[[0, 0, 0.5]], colors=[[0, 0, 255]]), static=True)
        except Exception:
            pass

        rec.log("world/model", rr.Asset3D(path=glb_path), static=True)

        if rrb is not None:
            try:
                blueprint = rrb.Blueprint(
                    rrb.Spatial3DView(
                        origin="/world",
                        name="3D View",
                    ),
                    collapse_panels=True,
                )
                rec.send_blueprint(blueprint)
            except Exception as e:
                print(f"Blueprint creation failed (non-fatal): {e}")

        rrd_path = os.path.join(user_dir, f'trellis_output_{timestamp}.rrd')
        rec.save(rrd_path)

        torch.cuda.empty_cache()
        return rrd_path, glb_path, glb_path, image_path

    except Exception as e:
        torch.cuda.empty_cache()
        raise gr.Error(f"Generation failed: {str(e)}")


# ====================================================================
# --- Video Generation (GLB to MP4) ---
# ====================================================================

def hex_to_rgb_norm(hex_color: str) -> tuple:
    """Convert hex color string to normalized (0-1) RGB tuple."""
    hex_color = hex_color.lstrip('#')
    r = int(hex_color[0:2], 16) / 255.0
    g = int(hex_color[2:4], 16) / 255.0
    b = int(hex_color[4:6], 16) / 255.0
    return (r, g, b)

def get_polygon_color_norm(color_name: str) -> tuple:
    """Convert color name to normalized RGB tuple."""
    colors_map = {
        "Pink": (1.0, 0.75, 0.80),
        "Green": (0.0, 1.0, 0.0),
        "White": (1.0, 1.0, 1.0),
        "Orange": (1.0, 0.65, 0.0)
    }
    return colors_map.get(color_name, (1.0, 0.75, 0.80))

def look_at_matrix(eye: np.ndarray, target: np.ndarray, up: np.ndarray = np.array([0, 1, 0])) -> np.ndarray:
    """Compute a 4x4 camera look-at matrix."""
    forward = target - eye
    forward = forward / (np.linalg.norm(forward) + 1e-8)
    right = np.cross(forward, up)
    right = right / (np.linalg.norm(right) + 1e-8)
    up_vec = np.cross(right, forward)

    pose = np.eye(4)
    pose[:3, 0] = right
    pose[:3, 1] = up_vec
    pose[:3, 2] = -forward
    pose[:3, 3] = eye
    return pose


def compute_camera_pose(style: str, t: float, rotation_speed: str = "Normal") -> np.ndarray:
    """
    Compute camera pose (4x4 matrix) for given style and normalized time t (0 to 1).
    The camera always looks at the origin where the model is centered.
    """
    speed_map = {"Slow": 0.5, "Normal": 1.0, "Fast": 2.0}
    speed = speed_map.get(rotation_speed, 1.0)

    angle = t * 2 * np.pi * speed
    base_distance = 3.0

    if style == "Orbit":
        x = base_distance * np.cos(angle)
        z = base_distance * np.sin(angle)
        y = 0.0
    elif style == "Zoom In":
        d = base_distance * (1.8 - t * 1.0)
        x = d * np.cos(angle)
        z = d * np.sin(angle)
        y = 0.0
    elif style == "Zoom Out":
        d = base_distance * (0.8 + t * 1.0)
        x = d * np.cos(angle)
        z = d * np.sin(angle)
        y = 0.0
    elif style == "Turntable":
        x = base_distance * np.cos(angle)
        z = base_distance * np.sin(angle)
        y = 0.3
    elif style == "Spiral":
        d = base_distance * (1.4 - t * 0.2)
        y = 2.0 * (1.0 - t)
        x = d * np.cos(angle * 1.5)
        z = d * np.sin(angle * 1.5)
    elif style == "Top Sweep":
        h = base_distance * np.sin(t * np.pi / 2.0)
        v = base_distance * np.cos(t * np.pi / 2.0) * 0.9
        x = h * np.cos(angle)
        z = h * np.sin(angle)
        y = v
    elif style == "Cinematic":
        d = base_distance * (1.0 + 0.3 * np.sin(t * 2 * np.pi))
        y = 0.5 * np.sin(t * 2 * np.pi)
        x = d * np.cos(angle)
        z = d * np.sin(angle)
    else:
        x = base_distance * np.cos(angle)
        z = base_distance * np.sin(angle)
        y = 0.0

    eye = np.array([x, y, z])
    target = np.array([0.0, 0.0, 0.0])
    return look_at_matrix(eye, target)


def compute_norm_transform(scene_or_mesh) -> np.ndarray:
    """Compute a 4x4 transform that centers and scales the object to fit in a unit sphere."""
    if isinstance(scene_or_mesh, trimesh.Scene):
        bounds = scene_or_mesh.bounds
    else:
        bounds = np.array([scene_or_mesh.bounds[0], scene_or_mesh.bounds[1]])

    center = (bounds[0] + bounds[1]) / 2.0
    extent = bounds[1] - bounds[0]
    scale = 2.0 / max(extent.max(), 1e-6)

    T = np.eye(4)
    T[:3, 3] = -center
    S = np.eye(4)
    S[:3, :3] *= scale
    return S @ T


def write_video_frames(frames: list, output_path: str, fps: int):
    """Write a list of RGB uint8 frames to an MP4 file using OpenCV."""
    if len(frames) == 0:
        raise ValueError("No frames to write.")

    height, width = frames[0].shape[:2]

    # Try H264/avc1 first, fall back to mp4v
    writer = None
    for codec in ['avc1', 'mp4v', 'X264']:
        fourcc = cv2.VideoWriter_fourcc(*codec)
        writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        if writer.isOpened():
            break
        writer = None

    if writer is None:
        raise RuntimeError("Could not open video writer with any codec.")

    for frame in frames:
        if frame.dtype != np.uint8:
            frame = (frame * 255).astype(np.uint8)
        if len(frame.shape) == 3 and frame.shape[2] == 3:
            frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        else:
            frame_bgr = frame
        writer.write(frame_bgr)

    writer.release()


def render_video_pyrender(
    glb_path: str,
    output_path: str,
    style: str,
    rotation_speed: str,
    bg_color_hex: str,
    polygon_color_name: str,
    resolution: tuple,
    fps: int = 30,
    duration: int = 5,
    progress: gr.Progress = None,
) -> str:
    """Render GLB to video using pyrender (high quality)."""
    tri_scene = trimesh.load(glb_path, force='scene')
    norm_transform = compute_norm_transform(tri_scene)

    bg_color = hex_to_rgb_norm(bg_color_hex)
    # pyrender expects RGBA for bg_color
    bg_color_rgba = (*bg_color, 1.0)

    pyr_scene = pyrender.Scene(bg_color=bg_color_rgba, ambient_light=[0.25, 0.25, 0.30])

    poly_color = get_polygon_color_norm(polygon_color_name)
    material = pyrender.MetallicRoughnessMaterial(
        baseColorFactor=[poly_color[0], poly_color[1], poly_color[2], 1.0],
        metallicFactor=0.1,
        roughnessFactor=0.5
    )

    # Add meshes
    if isinstance(tri_scene, trimesh.Scene):
        for node_name in tri_scene.graph.nodes_geometry:
            transform, geometry_name = tri_scene.graph.get(node_name)
            geometry = tri_scene.geometry[geometry_name]
            final_transform = norm_transform @ transform
            pr_mesh = pyrender.Mesh.from_trimesh(geometry, material=material, smooth=True)
            pyr_scene.add(pr_mesh, pose=final_transform)
    else:
        pr_mesh = pyrender.Mesh.from_trimesh(tri_scene, material=material, smooth=True)
        pyr_scene.add(pr_mesh, pose=norm_transform)

    # Add lights
    key_light = pyrender.DirectionalLight(color=[1.0, 1.0, 0.95], intensity=3.5)
    fill_light = pyrender.DirectionalLight(color=[0.6, 0.65, 0.8], intensity=1.5)
    rim_light = pyrender.DirectionalLight(color=[0.8, 0.8, 1.0], intensity=2.0)

    pyr_scene.add(key_light, pose=look_at_matrix(np.array([3, 3, 3]), np.array([0, 0, 0])))
    pyr_scene.add(fill_light, pose=look_at_matrix(np.array([-3, 1, -2]), np.array([0, 0, 0])))
    pyr_scene.add(rim_light, pose=look_at_matrix(np.array([0, 2, -4]), np.array([0, 0, 0])))

    # Camera
    camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0, aspectRatio=resolution[0] / resolution[1])
    camera_node = pyr_scene.add(camera, pose=np.eye(4))

    # Renderer
    renderer = pyrender.OffscreenRenderer(resolution[0], resolution[1])

    num_frames = fps * duration
    frames = []

    for i in range(num_frames):
        t = i / max(num_frames - 1, 1)
        cam_pose = compute_camera_pose(style, t, rotation_speed)
        pyr_scene.set_pose(camera_node, cam_pose)

        color, _ = renderer.render(pyr_scene)
        frames.append(color)

        if progress is not None:
            progress(0.1 + 0.85 * (i + 1) / num_frames, desc=f"Rendering frame {i+1}/{num_frames}...")

    renderer.delete()

    write_video_frames(frames, output_path, fps)
    return output_path


def render_video_matplotlib(
    glb_path: str,
    output_path: str,
    style: str,
    rotation_speed: str,
    bg_color_hex: str,
    polygon_color_name: str,
    resolution: tuple,
    fps: int = 30,
    duration: int = 5,
    progress: gr.Progress = None,
) -> str:
    """Render GLB to video using matplotlib (fallback, lower quality)."""
    tri_obj = trimesh.load(glb_path, force='mesh')

    # Simplify for performance
    if len(tri_obj.faces) > 8000:
        try:
            tri_obj = tri_obj.simplify_quadric_decimation(8000)
        except Exception:
            pass

    # Normalize
    vertices = tri_obj.vertices.copy().astype(np.float64)
    vertices -= vertices.mean(axis=0)
    max_extent = np.abs(vertices).max()
    if max_extent > 1e-6:
        vertices /= max_extent

    # Swap Y and Z axes to make the object upright in matplotlib's 3D coordinate system
    # Matplotlib uses Z as the vertical axis by default, while GLB typically uses Y.
    vertices = vertices[:, [0, 2, 1]]

    faces = tri_obj.faces

    # Get face colors
    poly_color = get_polygon_color_norm(polygon_color_name)
    face_colors = np.tile(np.array([poly_color]), (len(faces), 1))

    # Compute face normals for lighting
    v0 = vertices[faces[:, 0]]
    v1 = vertices[faces[:, 1]]
    v2 = vertices[faces[:, 2]]
    normals = np.cross(v1 - v0, v2 - v0)
    norm_lens = np.linalg.norm(normals, axis=1, keepdims=True)
    norm_lens[norm_lens < 1e-8] = 1.0
    normals /= norm_lens

    light_dir = np.array([0.5, 0.8, 1.0])
    light_dir = light_dir / np.linalg.norm(light_dir)
    shading = np.maximum(0, normals @ light_dir)
    shading = 0.35 + 0.65 * shading

    shaded_colors = face_colors * shading[:, np.newaxis]
    shaded_colors = np.clip(shaded_colors, 0, 1)

    bg_rgb = hex_to_rgb_norm(bg_color_hex)
    num_frames = fps * duration

    speed_map = {"Slow": 0.5, "Normal": 1.0, "Fast": 2.0}
    speed = speed_map.get(rotation_speed, 1.0)

    dpi = 100
    fig_w = resolution[0] / dpi
    fig_h = resolution[1] / dpi
    fig = plt.figure(figsize=(fig_w, fig_h), dpi=dpi)
    ax = fig.add_subplot(111, projection='3d')

    frames = []

    for i in range(num_frames):
        t = i / max(num_frames - 1, 1)
        angle = t * 360 * speed
        elev = 0  # Set to 0 for straight axis view

        zoom = 1.6
        if style == "Zoom In":
            zoom = 2.2 - t * 1.0
        elif style == "Zoom Out":
            zoom = 1.2 + t * 1.0
        elif style == "Spiral":
            elev = 45 - t * 45
            zoom = 1.8 - t * 0.3
        elif style == "Top Sweep":
            elev = 45 - t * 45
        elif style == "Turntable":
            elev = 0
        elif style == "Cinematic":
            elev = 10 + 10 * np.sin(t * 2 * np.pi)
            zoom = 1.6 + 0.2 * np.sin(t * 2 * np.pi)

        ax.clear()
        ax.set_facecolor(bg_rgb)
        fig.patch.set_facecolor(bg_rgb)

        poly = Poly3DCollection(vertices[faces], alpha=1.0)
        poly.set_facecolor(shaded_colors)
        poly.set_edgecolor('none')
        ax.add_collection3d(poly)

        ax.set_xlim(-zoom, zoom)
        ax.set_ylim(-zoom, zoom)
        ax.set_zlim(-zoom, zoom)
        ax.view_init(elev=elev, azim=angle)
        ax.set_axis_off()

        fig.canvas.draw()
        
        # FIX: Use buffer_rgba instead of tostring_rgb for newer matplotlib versions
        buf = np.asarray(fig.canvas.buffer_rgba())
        frame = buf[:, :, :3].copy()
        frames.append(frame)

        if progress is not None:
            progress(0.1 + 0.85 * (i + 1) / num_frames, desc=f"Rendering frame {i+1}/{num_frames} (matplotlib)...")

    plt.close(fig)

    write_video_frames(frames, output_path, fps)
    return output_path


def generate_video(
    glb_path: str,
    style: str,
    rotation_speed: str,
    bg_color: str,
    polygon_color: str,
    resolution_str: str,
    req: Optional[gr.Request] = None,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
) -> str:
    """Generate a 5-second MP4 video from a GLB file."""
    if glb_path is None or not os.path.exists(glb_path):
        raise gr.Error("No GLB file found. Please generate a 3D model first.")

    session_hash = req.session_hash if req is not None else "default_video_session"
    user_dir = os.path.join(TMP_DIR, session_hash)
    os.makedirs(user_dir, exist_ok=True)

    resolution_map = {
        "512x512": (512, 512),
        "720p (1280x720)": (1280, 720),
        "1080p (1920x1080)": (1920, 1080),
    }
    resolution = resolution_map.get(resolution_str, (512, 512))

    fps = 30
    duration = 5  # Fixed at 5 seconds as requested
    timestamp = datetime.now().strftime("%Y-%m-%dT%H%M%S")
    output_path = os.path.join(user_dir, f'video_{timestamp}.mp4')

    progress(0.05, desc="Loading 3D model...")

    try:
        if HAS_PYRENDER:
            progress(0.1, desc="Rendering with pyrender (high quality)...")
            try:
                return render_video_pyrender(
                    glb_path, output_path, style, rotation_speed, bg_color, polygon_color,
                    resolution, fps, duration, progress
                )
            except Exception as e:
                print(f"pyrender rendering failed: {e}. Falling back to matplotlib.")

        progress(0.1, desc="Rendering with matplotlib (fallback)...")
        return render_video_matplotlib(
            glb_path, output_path, style, rotation_speed, bg_color, polygon_color,
            resolution, fps, duration, progress
        )

    except Exception as e:
        raise gr.Error(f"Video generation failed: {str(e)}")


# ====================================================================
# --- ZIP Export ---
# ====================================================================

def random_name(length: int = 8) -> str:
    """Generate a random alphanumeric name."""
    return ''.join(random.choices(string.ascii_lowercase + string.digits, k=length))


def export_zip(
    image_path: str,
    glb_path: str,
    video_path: str,
    prompt: str,
    video_style: str,
    req: Optional[gr.Request] = None,
) -> str:
    """Export image, GLB, and video as a single zip file."""
    session_hash = req.session_hash if req is not None else "default_export_session"
    user_dir = os.path.join(TMP_DIR, session_hash)
    os.makedirs(user_dir, exist_ok=True)

    if (not image_path or not os.path.exists(image_path)) and \
       (not glb_path or not os.path.exists(glb_path)) and \
       (not video_path or not os.path.exists(video_path)):
        raise gr.Error("No files available to export. Please generate a 3D model and video first.")

    name = random_name()
    zip_path = os.path.join(user_dir, f'{name}.zip')

    metadata = {
        'prompt': prompt if prompt else '',
        'video_style': video_style if video_style else '',
        'timestamp': datetime.now().isoformat(),
    }

    try:
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
            if image_path and os.path.exists(image_path):
                ext = os.path.splitext(image_path)[1]
                zf.write(image_path, arcname=f'image{ext}')
            if glb_path and os.path.exists(glb_path):
                zf.write(glb_path, arcname='model.glb')
            if video_path and os.path.exists(video_path):
                ext = os.path.splitext(video_path)[1]
                zf.write(video_path, arcname=f'video{ext}')
            
            # Write metadata
            zf.writestr('metadata.json', json.dumps(metadata, indent=4))
    except Exception as e:
        raise gr.Error(f"Failed to create ZIP file: {str(e)}")

    return zip_path


# ====================================================================
# --- Gradio UI ---
# ====================================================================

css="""
#col-container {
    margin: 0 auto;
    max-width: 960px;
}
#main-title h1 {font-size: 2.4em !important;}
#video-section, #zip-section {
    margin-top: 20px;
    padding: 16px;
    border-radius: 12px;
    border: 2px solid var(--block-border-color);
}
"""

if __name__ == "__main__":
    os.makedirs(TMP_DIR, exist_ok=True)

    with gr.Blocks() as demo:
        gr.Markdown("# **Image-to-3D-Video-Asset-Generator**", elem_id="main-title")
        gr.Markdown("""
        Generate a complete 3D asset pipeline: Text-to-Image to 3D (GLB) to Video, then export everything as a ZIP file.
        Powered by [TRELLIS.2](https://huggingface.co/microsoft/TRELLIS.2-4B), [Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo), and custom rendering. [GitHub](https://github.com/PRITHIVSAKTHIUR/Image-to-3D-Video-Asset-Generator)
        """)

        # State variables for tracking file paths
        state_glb = gr.State(value=None)
        state_image = gr.State(value=None)
        state_video = gr.State(value=None)

        # --- Section 1: 3D Generation ---
        with gr.Row():
            with gr.Column(scale=1, min_width=360):

                with gr.Tabs():
                    with gr.Tab("Text-to-Image-3D"):
                        txt_prompt = gr.Textbox(label="Prompt", placeholder="eg. A Plane 3D model", lines=2)
                        btn_gen_img = gr.Button("1. Generate Image", variant="primary")
                    with gr.Tab("Image-to-3D"):
                        gr.Markdown("Upload an image directly if you have one.")

                image_prompt = gr.Image(label="Input Image", format="png", image_mode="RGBA", type="pil", height=350)

                with gr.Accordion(label="3D Settings", open=False):
                    resolution = gr.Radio(["512", "1024", "1536"], label="Generation Resolution", value="1024")
                    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(50000, 500000, label="Target Faces", value=150000, step=10000)
                texture_size = gr.Slider(512, 4096, label="Texture Size", value=1024, step=512)

                btn_gen_3d = gr.Button("2. Generate 3D", variant="primary", scale=2)

                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")
                    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")

            with gr.Column(scale=2):
                gr.Markdown("### 3D Output")

                rerun_output = Rerun(
                    label="Rerun 3D Viewer",
                    height=600
                )

                download_btn = gr.DownloadButton(label="3. Download GLB File", variant="primary")

                gr.Examples(
                    examples=[
                        ["example-images/A (1).webp"],
                        ["example-images/A (2).webp"],
                        ["example-images/A (3).webp"],
                        ["example-images/A (4).webp"],
                        ["example-images/A (5).webp"],
                        ["example-images/A (6).webp"],
                        ["example-images/A (7).webp"],
                        ["example-images/A (8).webp"],
                        ["example-images/A (9).webp"],
                        ["example-images/A (10).webp"],
                        ["example-images/A (11).webp"],
                        ["example-images/A (12).webp"],
                        ["example-images/A (13).webp"],
                        ["example-images/A (14).webp"],
                        ["example-images/A (15).webp"],
                        ["example-images/A (16).webp"],
                        ["example-images/A (17).webp"],
                        ["example-images/A (18).webp"],
                        ["example-images/A (19).webp"],
                        ["example-images/A (20).webp"],
                    ],
                    inputs=[image_prompt],
                    label="Image Examples [image-to-3d]"
                )

                gr.Examples(
                    examples=[
                        ["A Cat 3D model"],
                        ["A realistic Cat 3D model"],
                        ["A cartoon Cat 3D model"],
                        ["A low poly Cat 3D"],
                        ["A cyberpunk Cat 3D"],
                        ["A robotic Cat 3D"],
                        ["A fluffy Cat 3D"],
                        ["A fantasy Cat 3D creature"],
                        ["A stylized Cat 3D"],
                        ["A Cat 3D sculpture"],

                        ["A Plane 3D model"],
                        ["A commercial Plane 3D"],
                        ["A fighter jet Plane 3D"],
                        ["A low poly Plane 3D"],
                        ["A vintage Plane 3D"],
                        ["A futuristic Plane 3D"],
                        ["A cargo Plane 3D"],
                        ["A private jet Plane 3D"],
                        ["A toy Plane 3D"],
                        ["A realistic Plane 3D"],

                        ["A Car 3D model"],
                        ["A sports Car 3D"],
                        ["A luxury Car 3D"],
                        ["A low poly Car 3D"],
                        ["A racing Car 3D"],
                        ["A cyberpunk Car 3D"],
                        ["A vintage Car 3D"],
                        ["A futuristic Car 3D"],
                        ["A SUV Car 3D"],
                        ["A electric Car 3D"],

                        ["A Shoe 3D model"],
                        ["A sneaker Shoe 3D"],
                        ["A running Shoe 3D"],
                        ["A leather Shoe 3D"],
                        ["A high heel Shoe 3D"],
                        ["A boot Shoe 3D"],
                        ["A low poly Shoe 3D"],
                        ["A futuristic Shoe 3D"],
                        ["A sports Shoe 3D"],
                        ["A casual Shoe 3D"],

                        ["A Chair 3D model"],
                        ["A Table 3D model"],
                        ["A Sofa 3D model"],
                        ["A Lamp 3D model"],
                        ["A Watch 3D model"],
                        ["A Backpack 3D model"],
                        ["A Drone 3D model"],
                        ["A Robot 3D model"],
                        ["A Smartphone 3D model"],
                        ["A Headphones 3D model"],

                        ["A House 3D model"],
                        ["A Skyscraper 3D model"],
                        ["A Bridge 3D model"],
                        ["A Castle 3D model"],
                        ["A Spaceship 3D model"],
                        ["A Rocket 3D model"],
                        ["A Satellite 3D model"],
                        ["A Tank 3D model"],
                        ["A Motorcycle 3D model"],
                        ["A Bicycle 3D model"]
                    ],
                    inputs=[txt_prompt],
                    label="3D Prompt Examples [text-to-3d]"
                )

        # --- Section 2: 3D to Video Converter ---
        with gr.Column(elem_id="video-section"):
            gr.Markdown("## 3D (GLB) to Video Converter")
            gr.Markdown("Convert your generated 3D model into a 5-second MP4 video with customizable camera styles and effects.")

            with gr.Row():
                with gr.Column(scale=2):
                    video_style = gr.Dropdown(
                        choices=["Orbit", "Zoom In", "Zoom Out", "Turntable", "Spiral", "Top Sweep", "Cinematic"],
                        label="Video Style",
                        value="Zoom In",
                        info="Choose the camera movement style for the video."
                    )
                    rotation_speed = gr.Dropdown(
                        choices=["Slow", "Normal", "Fast"],
                        label="Rotation Speed",
                        value="Normal",
                        info="Controls how fast the camera rotates around the model."
                    )
                    polygon_color = gr.Dropdown(
                        choices=["Pink", "Green", "White", "Orange"],
                        label="Polygon Color",
                        value="Pink",
                        info="Override the 3D model's color in the video."
                    )
                with gr.Column(scale=2):
                    bg_color = gr.Dropdown(
                        choices=["#000000", "#FFFFFF", "#1A1A2E", "#2C3E50", "#0F0F0F", "#F0F0F0"],
                        label="Background Color",
                        value="#000000",
                        info="Background color for the video."
                    )
                    video_resolution = gr.Dropdown(
                        choices=["512x512", "720p (1280x720)", "1080p (1920x1080)"],
                        label="Video Resolution",
                        value="1080p (1920x1080)",
                        info="Higher resolutions take longer to render."
                    )

            with gr.Row():
                btn_gen_video = gr.Button("4. Generate Video (5 sec)", variant="primary")
                download_video_btn = gr.DownloadButton(label="5. Download Video", variant="secondary")

            video_output = gr.Video(label="Video Output", height=400)

        # --- Section 3: ZIP Export ---
        with gr.Column(elem_id="zip-section"):
            gr.Markdown("## ZIP Export")
            gr.Markdown("Export your generated assets (Image, 3D GLB, Video) as a single ZIP file.")

            with gr.Row():
                btn_export_zip = gr.Button("6. Generate ZIP File", variant="primary")
                download_zip_btn = gr.DownloadButton(label="7. Download ZIP File", variant="secondary")

        # --- Event Wiring ---
        demo.load(start_session)
        demo.unload(end_session)

        # Text-to-Image flow
        btn_gen_img.click(
            generate_txt2img,
            inputs=[txt_prompt],
            outputs=[image_prompt]
        ).then(
            preprocess_image,
            inputs=[image_prompt],
            outputs=[image_prompt]
        )

        image_prompt.upload(
            preprocess_image,
            inputs=[image_prompt],
            outputs=[image_prompt],
        )

        # 3D generation flow
        btn_gen_3d.click(
            get_seed,
            inputs=[randomize_seed, seed],
            outputs=[seed],
        ).then(
            generate_3d,
            inputs=[
                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=[rerun_output, download_btn, state_glb, state_image],
        )

        # Video generation flow
        def generate_video_wrapper(
            glb_path: str, 
            style: str, 
            rot_speed: str, 
            bg: str, 
            poly_color: str,
            res: str, 
            req: gr.Request, 
            progress: gr.Progress = gr.Progress(track_tqdm=True)
        ):
            video_path = generate_video(glb_path, style, rot_speed, bg, poly_color, res, req, progress)
            return video_path, video_path

        btn_gen_video.click(
            generate_video_wrapper,
            inputs=[state_glb, video_style, rotation_speed, bg_color, polygon_color, video_resolution],
            outputs=[video_output, state_video],
        ).then(
            lambda v: v,
            inputs=[state_video],
            outputs=[download_video_btn],
        )

        # ZIP export flow
        def export_zip_wrapper(
            image_path: str, 
            glb_path: str, 
            video_path: str, 
            prompt: str, 
            v_style: str, 
            req: gr.Request
        ):
            zip_path = export_zip(image_path, glb_path, video_path, prompt, v_style, req)
            return zip_path

        btn_export_zip.click(
            export_zip_wrapper,
            inputs=[state_image, state_glb, state_video, txt_prompt, video_style],
            outputs=[download_zip_btn],
        )

    demo.launch(theme=orange_red_theme, css=css, mcp_server=True, ssr_mode=False, show_error=True)