from functools import partial
import contextvars
import threading
import html
import re
import argparse
import base64
import inspect
import importlib.util
import json
import os
import sys
import time
from datetime import datetime
from pathlib import Path
from uuid import uuid4
from typing import Any, Dict, Generator, List, Optional, Tuple
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
RUNNING_ON_SPACES = bool(os.getenv("SPACE_ID"))
if RUNNING_ON_SPACES and importlib.util.find_spec("xformers") is not None:
_attn_backend = "xformers"
elif os.getenv("SPACE_ID"):
_attn_backend = "sdpa"
elif importlib.util.find_spec("flash_attn") is not None:
_attn_backend = "flash_attn"
elif importlib.util.find_spec("xformers") is not None:
_attn_backend = "xformers"
else:
_attn_backend = "sdpa"
os.environ.setdefault("ATTN_BACKEND", _attn_backend)
os.environ.setdefault("SPARSE_ATTN_BACKEND", _attn_backend)
os.environ.setdefault("SPARSE_ATTN", _attn_backend)
os.environ.setdefault("SPCONV_ALGO", "native")
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"
os.environ.setdefault("TRELLIS_REMBG_MODEL", "briaai/RMBG-2.0")
import cv2
import gradio as gr
import imageio
import imageio_ffmpeg
import numpy as np
import spaces
import logging
from bake_flight import BakeFlights
from quota_fallback import (
ALLOWANCE_ERROR,
CAPACITY_ERROR,
EXPORT,
GENERATE,
WAITING_STATUS,
FallbackCoordinator,
PendingCache,
is_eligible_embed,
is_quota_denial,
request_header,
safe_message,
)
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
FALLBACK = FallbackCoordinator(PendingCache(), {})
FLIGHTS = BakeFlights()
def configure_xformers_attention() -> None:
global _attn_backend
if not RUNNING_ON_SPACES or _attn_backend != "xformers":
return
try:
import xformers.ops as xformers_ops
original_memory_efficient_attention = xformers_ops.memory_efficient_attention
cutlass_attention_ops = (
xformers_ops.fmha.cutlass.FwOp,
xformers_ops.fmha.cutlass.BwOp,
)
except (ImportError, AttributeError, OSError) as error:
logger.warning("xformers unavailable; falling back to SDPA: %s", error)
_attn_backend = "sdpa"
os.environ["ATTN_BACKEND"] = _attn_backend
os.environ["SPARSE_ATTN_BACKEND"] = _attn_backend
os.environ["SPARSE_ATTN"] = _attn_backend
return
def memory_efficient_attention_with_cutlass(*args, **kwargs):
kwargs.setdefault("op", cutlass_attention_ops)
return original_memory_efficient_attention(*args, **kwargs)
xformers_ops.memory_efficient_attention = memory_efficient_attention_with_cutlass
configure_xformers_attention()
import torch
import trimesh
from PIL import Image
from gradio_litmodel3d import LitModel3D
from ground_slab import has_cast_shadow
from stl_export import DOWNLOAD_STEM, export_stl_from_state
from drive_storage import save_run_async
from usage_log import current_session_tag, gpu_call, labs_event, session_tag
from fallback_progress import current_event_id, mirror_progress
from access_limits import AccessPolicy, GrantBook, bridge_script, grants_required, public_allowance_from_env
import gradio_runtime
from starlette.templating import Jinja2Templates
sys.path.append(os.getcwd())
from trellis.pipelines import TrellisImageTo3DPipeline
from trellis.representations import Gaussian, MeshExtractResult
from trellis.utils import postprocessing_utils as trellis_postprocessing_utils
from trellis.utils import render_utils as trellis_render_utils
def patch_starlette_template_response() -> None:
"""Support Gradio's older TemplateResponse call style on newer Starlette."""
original_template_response = getattr(Jinja2Templates, "TemplateResponse", None)
if original_template_response is None:
logger.info("starlette TemplateResponse patch skipped: missing TemplateResponse")
return
if getattr(original_template_response, "__name__", "") == "patched_template_response_for_gradio":
return
try:
parameter_names = list(inspect.signature(original_template_response).parameters)
except Exception:
parameter_names = []
expects_request_first = len(parameter_names) > 1 and parameter_names[1] == "request"
if not expects_request_first:
logger.info("starlette TemplateResponse patch not needed")
return
def patched_template_response_for_gradio(self, *args, **kwargs):
if args and isinstance(args[0], str):
name = args[0]
context = args[1] if len(args) > 1 else kwargs.get("context")
if isinstance(context, dict) and "request" in context:
request = context["request"]
remaining_args = args[2:]
forwarded_kwargs = dict(kwargs)
forwarded_kwargs.pop("context", None)
return original_template_response(
self,
request,
name,
context,
*remaining_args,
**forwarded_kwargs,
)
return original_template_response(self, *args, **kwargs)
Jinja2Templates.TemplateResponse = patched_template_response_for_gradio
logger.info("starlette TemplateResponse patch applied")
patch_starlette_template_response()
parser = argparse.ArgumentParser(description="Pocket 3D AI")
parser.add_argument("--prod", action="store_true", help="Run in production mode")
parser.add_argument("--port", type=int, help="Port to run the server on (default: 8081 for prod, 8080 for dev)")
cmd_args, _unknown_args = parser.parse_known_args()
prod = cmd_args.prod
port = cmd_args.port if cmd_args.port else (8081 if prod else 8080)
show_options = not prod
PUBLIC_BASE_URL = os.getenv("PUBLIC_BASE_URL", "").strip()
if PUBLIC_BASE_URL and not PUBLIC_BASE_URL.startswith(("http://", "https://")):
PUBLIC_BASE_URL = f"https://{PUBLIC_BASE_URL}"
PUBLIC_BASE_URL = PUBLIC_BASE_URL.rstrip("/")
if RUNNING_ON_SPACES:
# Required behind proxies/custom domains so Gradio/Uvicorn generate https URLs.
os.environ.setdefault("FORWARDED_ALLOW_IPS", "*")
HAS_DIFF_GAUSSIAN_RASTERIZATION = importlib.util.find_spec("diff_gaussian_rasterization") is not None
if HAS_DIFF_GAUSSIAN_RASTERIZATION:
logger.info("diff_gaussian_rasterization detected: Gaussian RGB rendering enabled.")
else:
logger.warning(
"diff_gaussian_rasterization is missing: original TRELLIS textured preview/export is unavailable."
)
TRELLIS_RUNTIME = os.getenv("TRELLIS_RUNTIME", "original").strip().lower()
if TRELLIS_RUNTIME not in {"original", "trellis2"}:
logger.warning("Unknown TRELLIS_RUNTIME=%s, defaulting to 'original'.", TRELLIS_RUNTIME)
TRELLIS_RUNTIME = "original"
TRELLIS_SAMPLER = os.getenv("TRELLIS_SAMPLER", "default").strip().lower()
if TRELLIS_SAMPLER not in {"default", "mdt"}:
logger.warning("Unknown TRELLIS_SAMPLER=%s, defaulting to 'default'.", TRELLIS_SAMPLER)
TRELLIS_SAMPLER = "default"
TRELLIS_MDT_STEPS = int(os.getenv("TRELLIS_MDT_STEPS", "2"))
TRELLIS_BENCH = TRELLIS_RUNTIME == "original" and os.getenv("TRELLIS_BENCH", "").strip() == "1"
BENCH_GPU_SECONDS_PER_BAKE = 12
BENCH_MAX_GPU_SECONDS = 60
MAX_SEED = np.iinfo(np.int32).max
APP_DIR = os.path.dirname(os.path.abspath(__file__))
TMP_DIR = os.path.join(APP_DIR, "cache")
EXPORT_DIR = os.path.join(APP_DIR, "exports")
os.makedirs(TMP_DIR, exist_ok=True)
os.makedirs(EXPORT_DIR, exist_ok=True)
TMP_DIR_ABS = TMP_DIR
EXPORT_DIR_ABS = EXPORT_DIR
ASSETS_DIR_ABS = os.path.join(APP_DIR, "assets")
ALLOWED_PATHS = [TMP_DIR_ABS, EXPORT_DIR_ABS, ASSETS_DIR_ABS]
GPU_ALLOCATION_SECONDS = 18
GPU_ALLOCATION_SECONDS_STANDARD_TEXTURE = 12
STANDARD_TEXTURE_SIZE = 1024
GPU_ALLOCATION_SECONDS_HIGH_TEXTURE = 36
GPU_ALLOCATION_SECONDS_MAX_TEXTURE = 72
HIGH_TEXTURE_SIZE = 2048
MAX_TEXTURE_SIZE = 4096
GLB_EXPORT_SIMPLIFY = 0.95
GLB_EXPORT_TEXTURE_SIZE = 1024
GLB_TEXTURE_SIZES = (512, 1024, 1536, 2048, 4096)
GPU_PREVIEW, GPU_BAKE = "preview", "bake"
MULTI_IMAGE_MIN_VIEWS = 2
MULTI_IMAGE_MAX_VIEWS = 6
MULTI_IMAGE_EXAMPLE_DIR = os.path.join(ASSETS_DIR_ABS, "example_multi_image")
TRANSPARENT_PREVIEW = os.getenv("LABS_TRANSPARENT_PREVIEW", "").strip().lower() in {"1", "true", "yes"}
FAST_PREVIEW_RENDER = True
ALPHA_WEBM_SPEED_PARAMS = ()
GLB_MESH_KEY = "glb_mesh"
GLB_SETTINGS_KEY = "glb_settings"
GLB_REQUESTED_KEY = "glb_requested"
GLB_PAYLOAD_KEY = "glb_payload"
GLB_ERROR_KEY = "glb_error"
GLB_GENERATION_KEY = "glb_generation"
GLB_PREBAKED_KEY = "glb_prebaked"
GENERATED_VIA_KEY = "generated_via"
PREBAKE_GLB = True
GLB_PREP_FAILURE_MESSAGE = "Model generated, but texturing for download failed. Generate again to retry."
GLB_PACKAGING_FAILURE_MESSAGE = "Model generated, but packaging the download failed. Generate again to retry."
if gr.NO_RELOAD:
pipeline = None
envmap = None
trellis2_render_utils = None
EnvMap = None
o_voxel = None
mdt_flow_models = None
def load_distilled_flow_models():
global mdt_flow_models, TRELLIS_SAMPLER
try:
import mdt_flow
device = next(pipeline.models["sparse_structure_flow_model"].parameters()).device
mdt_flow_models = mdt_flow.load_flow_models(device)
logger.info("Loaded MDT-dist flow models on %s (sampler=%s, steps=%d).",
device, TRELLIS_SAMPLER, TRELLIS_MDT_STEPS)
except Exception as error:
logger.error("MDT-dist flow models failed to load; using the default sampler: %s", error, exc_info=True)
TRELLIS_SAMPLER = "default"
def initialize_pipeline():
global pipeline, envmap, trellis2_render_utils, EnvMap, o_voxel
if pipeline is not None:
return
logger.info("Initializing runtime '%s'...", TRELLIS_RUNTIME)
start_time = time.time()
try:
if TRELLIS_RUNTIME == "original":
pipeline = TrellisImageTo3DPipeline.from_pretrained(
"JeffreyXiang/TRELLIS-image-large",
formats=["mesh", "gaussian"],
)
if hasattr(pipeline, "_move_all_models_to_cpu"):
pipeline._move_all_models_to_cpu()
from trellis.pipelines.samplers.cached_flow_euler import CachedSlatSampler
pipeline.slat_sampler = CachedSlatSampler(
pipeline.slat_sampler.sigma_min, threshold=0.3, refine_final_steps=2,
)
logger.info("Using original TRELLIS runtime with conservative SLAT caching.")
if TRELLIS_SAMPLER == "mdt" or TRELLIS_BENCH:
load_distilled_flow_models()
pipeline.profile_stages = TRELLIS_BENCH
else:
from trellis2.pipelines import Trellis2ImageTo3DPipeline
from trellis2.renderers import EnvMap as Trellis2EnvMap
from trellis2.utils import render_utils as trellis2_render_utils_mod
import o_voxel as o_voxel_mod
pipeline = Trellis2ImageTo3DPipeline.from_pretrained("microsoft/TRELLIS.2-4B")
pipeline.low_vram = False
pipeline._device = "cpu"
trellis2_render_utils = trellis2_render_utils_mod
EnvMap = Trellis2EnvMap
o_voxel = o_voxel_mod
envmap = {}
for name in ["forest", "sunset", "courtyard"]:
exr_path = os.path.join("assets", "hdri", f"{name}.exr")
if os.path.exists(exr_path):
exr = cv2.imread(exr_path, cv2.IMREAD_UNCHANGED)
if exr is None:
continue
if RUNNING_ON_SPACES:
exr = cv2.resize(exr, (512, 256), interpolation=cv2.INTER_AREA)
envmap[name] = cv2.cvtColor(exr, cv2.COLOR_BGR2RGB)
logger.info("Using TRELLIS.2 runtime.")
logger.info("Pipeline initialized in %.2fs.", time.time() - start_time)
except Exception as e:
logger.error("Failed to initialize pipeline: %s", e, exc_info=True)
pipeline = None
raise
initialize_pipeline()
def clear_cuda_cache() -> None:
if torch.cuda.is_available():
torch.cuda.empty_cache()
def normalize_video_frames(frames: Any) -> List[np.ndarray]:
if frames is None:
return []
def _normalize_frame(arr: np.ndarray) -> Optional[np.ndarray]:
if arr is None or arr.ndim != 3:
return None
if arr.shape[-1] == 4:
arr = arr[:, :, :3]
if np.issubdtype(arr.dtype, np.floating):
min_v = float(np.nanmin(arr))
max_v = float(np.nanmax(arr))
if min_v >= 0.0 and max_v <= 1.0:
arr = arr * 255.0
elif min_v >= -1.0 and max_v <= 1.0:
arr = (arr + 1.0) * 127.5
arr = np.clip(arr, 0.0, 255.0)
return arr.astype(np.uint8)
if isinstance(frames, np.ndarray):
if frames.ndim == 4:
return [nf for f in frames if (nf := _normalize_frame(f)) is not None]
if frames.ndim == 3:
nf = _normalize_frame(frames)
return [nf] if nf is not None else []
return []
normalized = []
for frame in frames:
if frame is None:
continue
arr = np.asarray(frame)
nf = _normalize_frame(arr)
if nf is not None:
normalized.append(nf)
return normalized
def write_mp4(video_path: str, frames: List[np.ndarray], fps: int = 15) -> bool:
if not frames:
return False
try:
with imageio.get_writer(
video_path,
format="FFMPEG",
mode="I",
fps=fps,
codec="libx264",
ffmpeg_params=["-preset", "ultrafast", "-movflags", "faststart"],
) as writer:
for frame in frames:
writer.append_data(frame)
if os.path.exists(video_path) and os.path.getsize(video_path) > 0:
return True
except Exception as ffmpeg_err:
logger.warning("FFMPEG video writer failed: %s", ffmpeg_err)
try:
h, w = frames[0].shape[:2]
writer = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
if not writer.isOpened():
raise RuntimeError("OpenCV VideoWriter failed to open.")
for frame in frames:
if frame.shape[0] != h or frame.shape[1] != w:
frame = cv2.resize(frame, (w, h), interpolation=cv2.INTER_AREA)
writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
writer.release()
if os.path.exists(video_path) and os.path.getsize(video_path) > 0:
return True
except Exception as opencv_err:
logger.error("OpenCV video writer failed: %s", opencv_err, exc_info=True)
return False
def preview_video_suffix(frames: List[np.ndarray]) -> str:
if not TRANSPARENT_PREVIEW or not frames:
return "mp4"
first = frames[0]
if isinstance(first, np.ndarray) and first.ndim == 3 and first.shape[-1] == 4:
return "webm"
return "mp4"
def compose_alpha_frames(black: List[np.ndarray], white: List[np.ndarray]) -> List[np.ndarray]:
if len(black) != len(white):
return black
frames = []
for back, front in zip(black, white):
back_f = back.astype(np.float32) / 255.0
front_f = front.astype(np.float32) / 255.0
alpha = np.clip(1.0 - (front_f - back_f).mean(axis=2, keepdims=True), 0.0, 1.0)
safe_alpha = np.where(alpha > 1e-4, alpha, 1.0)
rgb = np.clip(back_f / safe_alpha, 0.0, 1.0)
frame = np.concatenate([rgb * 255.0, alpha * 255.0], axis=2)
frames.append(np.rint(frame).astype(np.uint8))
return frames
def write_alpha_webm(video_path: str, frames: List[np.ndarray], fps: int = 15) -> bool:
if not frames:
return False
height, width = frames[0].shape[:2]
writer = None
try:
writer = imageio_ffmpeg.write_frames(
video_path,
(width, height),
pix_fmt_in="rgba",
pix_fmt_out="yuva420p",
fps=fps,
quality=None,
codec="libvpx",
ffmpeg_log_level="quiet",
output_params=[
"-b:v", "0", "-crf", "30", "-qmin", "0", "-qmax", "50", "-auto-alt-ref", "0",
*ALPHA_WEBM_SPEED_PARAMS,
],
)
writer.send(None)
for frame in frames:
writer.send(np.ascontiguousarray(frame))
except Exception as alpha_err:
logger.warning("Alpha WebM writer failed: %s", alpha_err)
return False
finally:
if writer is not None:
try:
writer.close()
except Exception as close_err:
logger.warning("Alpha WebM writer close failed: %s", close_err)
return os.path.exists(video_path) and os.path.getsize(video_path) > 0
def write_preview_video(frames: List[np.ndarray], fps: int = 15) -> Optional[str]:
transparent = preview_video_suffix(frames) == "webm"
if transparent:
alpha_path = os.path.join(TMP_DIR, f"{uuid4().hex}.webm")
if write_alpha_webm(alpha_path, frames, fps=fps):
return alpha_path
logger.warning("Alpha preview encoding failed; falling back to the opaque H.264 preview.")
opaque_frames = [frame[:, :, :3] for frame in frames] if transparent else frames
video_path = os.path.join(TMP_DIR, f"{uuid4().hex}.mp4")
return video_path if write_mp4(video_path, opaque_frames, fps=fps) else None
def is_video_mostly_black(frames: List[np.ndarray], mean_threshold: float = 2.0) -> bool:
if not frames:
return True
sample = frames[: min(8, len(frames))]
means = [float(np.mean(f)) for f in sample if f is not None]
if not means:
return True
return (sum(means) / len(means)) < mean_threshold
def start_session(req: gr.Request):
clear_cuda_cache()
os.makedirs(TMP_DIR, exist_ok=True)
os.makedirs(EXPORT_DIR, exist_ok=True)
def end_session(req: gr.Request):
FALLBACK.invalidate_session(session_hash(req))
clear_cuda_cache()
def session_hash(req: Optional[gr.Request]) -> str:
return getattr(req, "session_hash", "") or ""
def preprocess_image(image: Optional[Image.Image]) -> Optional[Image.Image]:
if image is None:
return None
try:
return pipeline.preprocess_image(image)
except Exception as e:
logger.error("Error during image preprocessing: %s", e, exc_info=True)
return None
def preprocess_images(images: List[Image.Image]) -> Optional[List[Image.Image]]:
if not images:
return None
try:
processed = pipeline.preprocess_images(images)
if any(img is None for img in processed):
logger.error("One or more images failed preprocessing.")
return None
return processed
except Exception as e:
logger.error("Error during multi-image preprocessing: %s", e, exc_info=True)
return None
def normalize_gallery_images(value: Any) -> List[Image.Image]:
images: List[Image.Image] = []
for item in value or []:
if isinstance(item, (tuple, list)) and item:
item = item[0]
if isinstance(item, Image.Image):
images.append(item)
elif isinstance(item, str):
images.append(Image.open(item))
return images
def validate_multi_images(value: Any) -> List[Image.Image]:
images = normalize_gallery_images(value)
if len(images) < MULTI_IMAGE_MIN_VIEWS:
raise gr.Error(f"Add at least {MULTI_IMAGE_MIN_VIEWS} images to generate from multiple views.")
if len(images) > MULTI_IMAGE_MAX_VIEWS:
raise gr.Error(f"Use at most {MULTI_IMAGE_MAX_VIEWS} images for one generation.")
return images
def prepare_multi_example() -> List[Image.Image]:
if not os.path.isdir(MULTI_IMAGE_EXAMPLE_DIR):
return []
cases = sorted({name.split("_")[0] for name in os.listdir(MULTI_IMAGE_EXAMPLE_DIR) if name.endswith(".png")})
concatenated: List[Image.Image] = []
for case in cases:
views = []
for index in range(1, 4):
path = os.path.join(MULTI_IMAGE_EXAMPLE_DIR, f"{case}_{index}.png")
if not os.path.exists(path):
views = []
break
view = Image.open(path)
width, height = view.size
views.append(np.array(view.resize((int(width / height * 512), 512))))
if views:
concatenated.append(Image.fromarray(np.concatenate(views, axis=1)))
return concatenated
def split_multi_example(image: Image.Image) -> List[Image.Image]:
array = np.array(image)
opaque = np.any(array[..., 3] > 0, axis=0)
starts = np.where(~opaque[:-1] & opaque[1:])[0].tolist()
ends = np.where(opaque[:-1] & ~opaque[1:])[0].tolist()
views = [Image.fromarray(array[:, start:end + 1]) for start, end in zip(starts, ends)]
return views or [image]
def get_seed(randomize_seed: bool, seed: int) -> int:
return np.random.randint(0, MAX_SEED) if randomize_seed else seed
def encode_file_to_b64(file_path: str) -> str:
with open(file_path, "rb") as f:
return base64.b64encode(f.read()).decode("ascii")
def decode_b64_to_file(file_b64: str, out_path: str) -> None:
with open(out_path, "wb") as f:
f.write(base64.b64decode(file_b64))
def build_export_payload(glb_path: str) -> Dict[str, Any]:
return {
"glb_name": os.path.basename(glb_path),
"glb_b64": encode_file_to_b64(glb_path),
}
def get_public_base_url(req: Optional[gr.Request]) -> str:
if PUBLIC_BASE_URL:
return PUBLIC_BASE_URL
if req is not None:
headers = req.headers or {}
x_forwarded_host = headers.get("x-forwarded-host", "").split(",")[0].strip()
host = x_forwarded_host or headers.get("host", "").split(",")[0].strip()
if host:
return f"https://{host}"
space_host = os.getenv("SPACE_HOST", "").strip()
if space_host:
return f"https://{space_host}"
return ""
def to_public_file_value(file_path: str, req: Optional[gr.Request]) -> str:
base_url = get_public_base_url(req)
if RUNNING_ON_SPACES and base_url:
return f"{base_url}/file={file_path}"
return file_path
def materialize_export_payload(
payload: Optional[Dict[str, Any]],
req: gr.Request,
) -> Tuple[Dict[str, Any], Optional[str]]:
if payload is None:
return gr.update(value=None, visible=False), None
try:
download_dir = os.path.join(EXPORT_DIR, uuid4().hex)
os.makedirs(download_dir, exist_ok=True)
glb_path_abs = os.path.abspath(os.path.join(download_dir, f"{DOWNLOAD_STEM}.glb"))
decode_b64_to_file(str(payload["glb_b64"]), glb_path_abs)
glb_size = os.path.getsize(glb_path_abs)
if glb_size <= 0:
raise RuntimeError("Materialized GLB file is empty.")
logger.info("Materialized GLB for UI: %s (%d bytes)", glb_path_abs, glb_size)
glb_value = glb_path_abs
return gr.update(value=glb_value, visible=True), glb_value
except Exception as e:
logger.error("Materializing export payload failed: %s", e, exc_info=True)
raise gr.Error("Model generated, but packaging for display failed. Please retry.") from e
def generate_stl_download(state, cached_stl_path=None):
if not state:
raise gr.Error("Generate a model before downloading STL.")
if cached_stl_path and os.path.isfile(cached_stl_path):
return gr.update(value=cached_stl_path, label="Download STL", interactive=True), cached_stl_path
started, session, path = time.perf_counter(), current_session_tag(), None
try:
path, watertight = export_stl_from_state(state, EXPORT_DIR)
except ValueError as error:
raise gr.Error(str(error)) from error
finally:
labs_event("stl_done", session, ok=path is not None, took_s=round(time.perf_counter() - started, 2))
if not watertight:
gr.Warning("This STL has open edges. Check and repair it in your slicer before printing.")
logger.info("CPU STL export completed in %.2fs; watertight=%s", time.perf_counter() - started, watertight)
return gr.update(value=path, label="Download STL", interactive=True), path
def reset_generation_outputs():
return (
None, None, gr.update(value=None, visible=False), None, None, None,
gr.update(value=None, label="Generate STL", interactive=False),
gr.update(interactive=False, visible=True), gr.update(value="", visible=False),
gr.update(value=None, visible=False),
)
def begin_generation(randomize_seed, seed):
return (get_seed(randomize_seed, seed), *reset_generation_outputs())
def requested_glb_settings(state: Optional[Dict[str, Any]]) -> Dict[str, Any]:
settings = (state or {}).get(GLB_REQUESTED_KEY) or (state or {}).get(GLB_SETTINGS_KEY) or {}
return {
"simplify": float(settings.get("simplify", GLB_EXPORT_SIMPLIFY)),
"texture_size": int(settings.get("texture_size", GLB_EXPORT_TEXTURE_SIZE)),
}
def enable_model_exports(state):
ready = state is not None
return (
gr.update(value=None, label="Generate STL", interactive=ready),
gr.update(interactive=ready, visible=True), None,
requested_glb_settings(state)["texture_size"],
)
def begin_glb_export():
return gr.update(interactive=False), gr.update(value="Preparing download...", visible=True)
def finish_glb_export(glb_path, state, status_text):
failed = glb_path is None
message = (status_text or "").strip() or "Export failed. Please try again."
return (
gr.update(interactive=state is not None and failed, visible=failed),
gr.update(value=message if failed else "", visible=failed),
gr.update(value=None if failed else glb_path, visible=not failed),
)
def get_preview_settings(req: Optional[gr.Request]) -> Tuple[bool, int, int, int]:
user_agent = request_header(req, "user-agent").lower()
is_mobile = any(d in user_agent for d in ["android", "iphone", "ipad", "mobile"])
resolution = 256 if is_mobile else 384
fps = 15
seconds = 3
return is_mobile, resolution, fps, seconds
def pack_original_state(outputs: Dict[str, Any]) -> Dict[str, Any]:
mesh = outputs["mesh"][0]
packed: Dict[str, Any] = {
"runtime": "original",
"mesh": {
"vertices": mesh.vertices.detach().cpu().numpy(),
"faces": mesh.faces.detach().cpu().numpy(),
"vertex_attrs": mesh.vertex_attrs.detach().cpu().numpy() if mesh.vertex_attrs is not None else None,
"res": int(getattr(mesh, "res", 64)),
},
}
gaussian_outputs = outputs.get("gaussian", None)
if gaussian_outputs and gaussian_outputs[0] is not None:
gaussian = gaussian_outputs[0]
packed["gaussian"] = {
"init_params": dict(getattr(gaussian, "init_params", {})),
"xyz": gaussian._xyz.detach().cpu().numpy(),
"features_dc": gaussian._features_dc.detach().cpu().numpy(),
"features_rest": gaussian._features_rest.detach().cpu().numpy() if gaussian._features_rest is not None else None,
"scaling": gaussian._scaling.detach().cpu().numpy(),
"rotation": gaussian._rotation.detach().cpu().numpy(),
"opacity": gaussian._opacity.detach().cpu().numpy(),
}
return packed
def unpack_original_state(state: Dict[str, Any]) -> Dict[str, Any]:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
minfo = state["mesh"]
vertex_attrs_np = minfo.get("vertex_attrs")
mesh = MeshExtractResult(
vertices=torch.tensor(minfo["vertices"], device=device, dtype=torch.float32),
faces=torch.tensor(minfo["faces"], device=device, dtype=torch.int64),
vertex_attrs=torch.tensor(vertex_attrs_np, device=device, dtype=torch.float32) if vertex_attrs_np is not None else None,
res=int(minfo.get("res", 64)),
)
out: Dict[str, Any] = {"mesh": mesh}
ginfo = state.get("gaussian")
if ginfo is not None:
gaussian = Gaussian(device=device.type, **ginfo["init_params"])
gaussian._xyz = torch.tensor(ginfo["xyz"], device=device, dtype=torch.float32)
gaussian._features_dc = torch.tensor(ginfo["features_dc"], device=device, dtype=torch.float32)
gaussian._features_rest = (
torch.tensor(ginfo["features_rest"], device=device, dtype=torch.float32)
if ginfo["features_rest"] is not None else None
)
gaussian._scaling = torch.tensor(ginfo["scaling"], device=device, dtype=torch.float32)
gaussian._rotation = torch.tensor(ginfo["rotation"], device=device, dtype=torch.float32)
gaussian._opacity = torch.tensor(ginfo["opacity"], device=device, dtype=torch.float32)
out["gaussian"] = gaussian
return out
def pack_trellis2_state(mesh: Any, grid_size: int) -> Dict[str, Any]:
return {
"runtime": "trellis2",
"mesh": {
"vertices": mesh.vertices.detach().cpu().numpy(),
"faces": mesh.faces.detach().cpu().numpy(),
"attrs": mesh.attrs.detach().cpu().numpy(),
"coords": mesh.coords.detach().cpu().numpy(),
"voxel_shape": list(mesh.voxel_shape),
"layout": {k: [v.start, v.stop] for k, v in mesh.layout.items()},
},
"grid_size": grid_size,
}
def unpack_trellis2_state(state: Dict[str, Any]) -> Dict[str, Any]:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
mesh_info = state["mesh"]
attr_layout = {k: slice(v[0], v[1]) for k, v in mesh_info["layout"].items()}
return {
"vertices": torch.tensor(mesh_info["vertices"], device=device, dtype=torch.float32),
"faces": torch.tensor(mesh_info["faces"], device=device, dtype=torch.int32),
"attrs": torch.tensor(mesh_info["attrs"], device=device, dtype=torch.float32),
"coords": torch.tensor(mesh_info["coords"], device=device, dtype=torch.int32),
"voxel_shape": torch.Size(mesh_info["voxel_shape"]),
"attr_layout": attr_layout,
"grid_size": int(state["grid_size"]),
}
def render_gaussian_frames(
gaussian: Any, resolution: int, num_frames: int, bg_color: Tuple[float, float, float],
) -> List[np.ndarray]:
result = trellis_render_utils.render_video(
gaussian, resolution=resolution, num_frames=num_frames, bg_color=bg_color,
mode="color", verbose=False,
)
return result.get("color", []) if result else []
def compose_alpha_frames_on_gpu(black: Any, white: Any) -> Any:
back, front = black.float() / 255.0, white.float() / 255.0
alpha = (1.0 - (front - back).mean(dim=-1, keepdim=True)).clamp(0.0, 1.0)
rgb = (back / torch.where(alpha > 1e-4, alpha, torch.ones_like(alpha))).clamp(0.0, 1.0)
return torch.round(torch.cat([rgb * 255.0, alpha * 255.0], dim=-1)).to(torch.uint8)
def render_preview_frames_on_gpu(gaussian: Any, resolution: int, frame_count: int) -> List[np.ndarray]:
backgrounds = [(0, 0, 0), (1, 1, 1)] if TRANSPARENT_PREVIEW else [(0, 0, 0)]
frames = trellis_render_utils.render_gaussian_turntable(gaussian, resolution, frame_count, backgrounds)
if float(frames[0, :8].float().mean()) < 2.0:
raise RuntimeError("RGB preview rendering produced blank color frames.")
if TRANSPARENT_PREVIEW:
frames = compose_alpha_frames_on_gpu(frames[0], frames[1])[None]
return list(frames[0].cpu().numpy())
def render_original_preview_frames(outputs: Dict[str, Any], req: gr.Request) -> Optional[List[np.ndarray]]:
_is_mobile, resolution, fps, seconds = get_preview_settings(req)
gaussian_outputs = outputs.get("gaussian", None)
gaussian = gaussian_outputs[0] if gaussian_outputs and gaussian_outputs[0] is not None else None
if gaussian is None:
raise RuntimeError("Original TRELLIS preview requires Gaussian appearance output.")
if not HAS_DIFF_GAUSSIAN_RASTERIZATION:
raise RuntimeError("diff_gaussian_rasterization is required for RGB preview rendering.")
frame_count = seconds * fps
if FAST_PREVIEW_RENDER:
return render_preview_frames_on_gpu(gaussian, resolution, frame_count)
frames = render_gaussian_frames(gaussian, resolution, frame_count, (0, 0, 0))
if not frames:
logger.warning("Preview rendering returned no frames.")
return None
if is_video_mostly_black(frames):
raise RuntimeError("RGB preview rendering produced blank color frames.")
if not TRANSPARENT_PREVIEW:
return frames
white_frames = render_gaussian_frames(gaussian, resolution, frame_count, (1, 1, 1))
if len(white_frames) != len(frames):
logger.warning(
"Alpha matte pass returned %d of %d frames; preview stays opaque.",
len(white_frames), len(frames),
)
return frames
return compose_alpha_frames(frames, white_frames)
def render_trellis2_preview_frames(mesh: Any, req: gr.Request) -> Optional[List[np.ndarray]]:
_is_mobile, resolution, fps, seconds = get_preview_settings(req)
num_frames = seconds * fps
loaded_envmap = {}
for name, exr_data in (envmap or {}).items():
loaded_envmap[name] = EnvMap(torch.tensor(exr_data, dtype=torch.float32, device="cuda"))
preview_envmap = loaded_envmap.get("sunset") if loaded_envmap else None
if preview_envmap is None and loaded_envmap:
preview_envmap = next(iter(loaded_envmap.values()))
if preview_envmap is not None:
vid_result = trellis2_render_utils.render_video(
mesh,
resolution=resolution,
num_frames=num_frames,
r=2,
fov=36,
envmap=preview_envmap,
)
else:
vid_result = trellis2_render_utils.render_video(
mesh,
resolution=resolution,
num_frames=num_frames,
r=2,
fov=36,
envmap=loaded_envmap,
)
shaded_frames = vid_result.get("shaded")
if shaded_frames is None:
shaded_keys = [k for k in vid_result.keys() if k.startswith("shaded_")]
if shaded_keys:
shaded_frames = vid_result[shaded_keys[0]]
color_frames = normalize_video_frames(shaded_frames if shaded_frames is not None else vid_result.get("color", []))
normal_frames = normalize_video_frames(vid_result.get("normal", []))
if len(color_frames) == 0 and len(normal_frames) == 0:
return None
return color_frames if len(color_frames) > 0 else normal_frames
def _run_original_pipeline(
image: Image.Image,
seed: int,
ss_guidance_strength: float,
ss_sampling_steps: int,
slat_guidance_strength: float,
slat_sampling_steps: int,
sampler: Optional[Dict[str, Any]] = None,
strip_ground_slab: bool = False,
) -> Dict[str, Any]:
sampler = sampler or {"name": TRELLIS_SAMPLER}
if sampler["name"] == "mdt":
return _run_distilled_pipeline(image, seed, int(sampler.get("steps", TRELLIS_MDT_STEPS)))
return pipeline.run(
image,
seed=seed,
formats=["mesh", "gaussian"],
preprocess_image=False,
sparse_structure_sampler_params=sparse_structure_params(ss_sampling_steps, ss_guidance_strength, strip_ground_slab),
slat_sampler_params={
"steps": slat_sampling_steps,
"cfg_strength": slat_guidance_strength,
},
)
def sparse_structure_params(steps, cfg_strength, strip_ground_slab):
return {"steps": steps, "cfg_strength": cfg_strength, "strip_ground_slab": strip_ground_slab}
def _run_distilled_pipeline(image: Image.Image, seed: int, steps: int) -> Dict[str, Any]:
import mdt_flow
if mdt_flow_models is None:
raise RuntimeError("MDT-dist flow models are not loaded.")
ss_params, slat_params = mdt_flow.sampler_params(steps)
with mdt_flow.distilled_sampling(pipeline, mdt_flow_models):
return pipeline.run(
image,
seed=seed,
formats=["mesh", "gaussian"],
preprocess_image=False,
sparse_structure_sampler_params=ss_params,
slat_sampler_params=slat_params,
)
def _run_original_pipeline_multi(
images: List[Image.Image],
seed: int,
ss_guidance_strength: float,
ss_sampling_steps: int,
slat_guidance_strength: float,
slat_sampling_steps: int,
strip_ground_slab: bool = False,
) -> Dict[str, Any]:
return pipeline.run_multi_image(
images,
seed=seed,
formats=["mesh", "gaussian"],
preprocess_image=False,
mode="stochastic",
sparse_structure_sampler_params=sparse_structure_params(ss_sampling_steps, ss_guidance_strength, strip_ground_slab),
slat_sampler_params={
"steps": slat_sampling_steps,
"cfg_strength": slat_guidance_strength,
},
)
def _run_trellis2_pipeline(
image: Image.Image,
seed: int,
ss_guidance_strength: float,
ss_sampling_steps: int,
slat_guidance_strength: float,
slat_sampling_steps: int,
) -> Tuple[Any, int]:
if o_voxel is None:
raise RuntimeError("TRELLIS.2 runtime requires o_voxel, but it is not available.")
pipeline.cuda()
if RUNNING_ON_SPACES:
pipeline_type = "512"
grid_size = 512
else:
pipeline_type = "1024_cascade"
grid_size = 1024
outputs = pipeline.run(
image,
seed=seed,
preprocess_image=False,
sparse_structure_sampler_params={
"steps": ss_sampling_steps,
"guidance_strength": ss_guidance_strength,
},
shape_slat_sampler_params={
"steps": slat_sampling_steps,
"guidance_strength": slat_guidance_strength,
},
tex_slat_sampler_params={
"steps": slat_sampling_steps,
"guidance_strength": slat_guidance_strength,
},
pipeline_type=pipeline_type,
return_latent=False,
)
return outputs[0], grid_size
def _prepare_glb(state: Dict[str, Any], texture_size: int) -> Any:
if state.get("runtime") == "original":
original_state = unpack_original_state(state)
mesh = original_state["mesh"]
app_rep = original_state.get("gaussian")
if app_rep is None:
raise RuntimeError("Original TRELLIS export requires Gaussian appearance output.")
if not HAS_DIFF_GAUSSIAN_RASTERIZATION:
raise RuntimeError("diff_gaussian_rasterization is required for textured GLB export.")
from export_profiles import balanced_export_settings, simplify_ratio_for_texture_size
simplify = simplify_ratio_for_texture_size(texture_size, len(mesh.faces))
return trellis_postprocessing_utils.to_glb(
app_rep, mesh, verbose=False,
**balanced_export_settings(simplify, texture_size),
)
if o_voxel is None:
raise gr.Error("TRELLIS.2 runtime requires o_voxel, but it is not available.")
from export_profiles import faces_for_texture_size
mesh_state = unpack_trellis2_state(state)
return o_voxel.postprocess.to_glb(
vertices=mesh_state["vertices"],
faces=mesh_state["faces"],
attr_volume=mesh_state["attrs"],
coords=mesh_state["coords"],
attr_layout=mesh_state["attr_layout"],
grid_size=mesh_state["grid_size"],
aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
decimation_target=faces_for_texture_size(texture_size),
texture_size=texture_size,
remesh=not RUNNING_ON_SPACES,
remesh_band=1,
remesh_project=0,
use_tqdm=False,
)
def save_export_payload(glb, progress):
if progress is not None:
progress(0.84, desc="Packaging GLB")
current_time_glb = datetime.now().strftime("%Y-%m%d-%H%M%S-%f")
glb_path = os.path.join(EXPORT_DIR, f"{current_time_glb}.glb")
glb.export(glb_path)
glb_path_abs = os.path.abspath(glb_path)
logger.info("GLB exported: %s (%d bytes)", glb_path_abs, os.path.getsize(glb_path_abs))
try:
scene_or_mesh = trimesh.load(glb_path_abs, force="scene")
if isinstance(scene_or_mesh, trimesh.Scene):
if not scene_or_mesh.geometry:
raise RuntimeError("Exported GLB contains no geometry.")
elif isinstance(scene_or_mesh, trimesh.Trimesh):
if scene_or_mesh.vertices is None or len(scene_or_mesh.vertices) == 0:
raise RuntimeError("Exported GLB mesh is empty.")
except Exception as validate_err:
logger.error("GLB validation failed: %s", validate_err, exc_info=True)
raise gr.Error("Model was exported but failed validation for viewer rendering.")
logger.info("Returning model file: %s", glb_path_abs)
return build_export_payload(glb_path_abs)
def bake_glb_on_gpu(state, mesh_simplify: float, texture_size: int, progress=None):
started = time.perf_counter()
torch.cuda.reset_peak_memory_stats()
try:
glb = _prepare_glb(state, int(texture_size))
except Exception as error:
logger.error("GLB preparation failed inside the GPU allocation: %s", error, exc_info=True)
message = plain_error_text(error) if isinstance(error, gr.Error) else GLB_PREP_FAILURE_MESSAGE
return None, message
finally:
clear_cuda_cache()
logger.info("GLB Bake Time: %.2fs (texture_size=%s)", time.perf_counter() - started, texture_size)
logger.info("GPU Peak Allocated: %.2fGiB", torch.cuda.max_memory_allocated() / (1024 ** 3))
return glb, None
def cache_prepared_glb(state: Optional[Dict[str, Any]], progress=None) -> Optional[Dict[str, Any]]:
if not state or GLB_MESH_KEY not in state:
return state
glb = state.pop(GLB_MESH_KEY)
started = time.perf_counter()
try:
state[GLB_PAYLOAD_KEY] = save_export_payload(glb, progress)
except Exception as error:
logger.error("GLB packaging failed after the GPU allocation: %s", error, exc_info=True)
state.pop(GLB_PAYLOAD_KEY, None)
state[GLB_ERROR_KEY] = (
plain_error_text(error) if isinstance(error, gr.Error) else GLB_PACKAGING_FAILURE_MESSAGE
)
return state
state.pop(GLB_ERROR_KEY, None)
logger.info("GLB Packaging Time: %.2fs", time.perf_counter() - started)
return state
def cached_glb_payload(
state: Optional[Dict[str, Any]], mesh_simplify: float, texture_size: int,
) -> Optional[Dict[str, Any]]:
if not state:
return None
settings = state.get(GLB_SETTINGS_KEY) or {}
if int(settings.get("texture_size", 0)) != int(texture_size):
return None
if abs(float(settings.get("simplify", -1.0)) - float(mesh_simplify)) > 1e-6:
return None
return state.get(GLB_PAYLOAD_KEY)
def _generate_and_render_model(
args,
req,
progress,
job_start,
texture_size=GLB_EXPORT_TEXTURE_SIZE,
original_runner=_run_original_pipeline,
):
outputs = None
mesh = None
grid_size = 512
inference_start = time.time()
progress(0.15, desc="Generating 3D structure")
if TRELLIS_RUNTIME == "original":
outputs = original_runner(*args)
else:
mesh, grid_size = _run_trellis2_pipeline(*args)
logger.info("Inference Time: %.2fs", time.time() - inference_start)
preview_start = time.time()
progress(0.45, desc="Rendering preview")
if TRELLIS_RUNTIME == "original":
state = pack_original_state(outputs)
frames = render_original_preview_frames(outputs, req)
else:
state = pack_trellis2_state(mesh, grid_size)
frames = render_trellis2_preview_frames(mesh, req)
logger.info("Preview Render Time: %.2fs", time.time() - preview_start)
state[GLB_REQUESTED_KEY] = {"simplify": GLB_EXPORT_SIMPLIFY, "texture_size": int(texture_size)}
progress(0.9, desc="Finishing preview")
logger.info("Total GPU Time: %.2fs", time.time() - job_start)
logger.info("GPU Peak Allocated: %.2fGiB", torch.cuda.max_memory_allocated() / (1024 ** 3))
if TRELLIS_BENCH and TRELLIS_RUNTIME == "original":
state["bench_timings"] = {
**getattr(pipeline, "stage_timings", {}),
"inference": round(preview_start - inference_start, 3),
"preview": round(time.time() - preview_start, 3),
"generate_total": round(time.time() - job_start, 3),
}
logger.info("STAGE_TIMINGS %s", json.dumps(state["bench_timings"]))
return frames, state
def _prepare_generation_input(image):
if not isinstance(image, (list, tuple)):
prepared = preprocess_image(image)
if prepared is None:
raise gr.Error("Image preprocessing failed. Please try a different image.")
return prepared, partial(_run_original_pipeline, strip_ground_slab=has_cast_shadow(image)), 1
views = normalize_gallery_images(image)
if TRELLIS_RUNTIME != "original":
raise gr.Error("Multiple-image generation requires the original TRELLIS runtime.")
if len(views) < MULTI_IMAGE_MIN_VIEWS:
raise gr.Error(f"Add at least {MULTI_IMAGE_MIN_VIEWS} images to generate from multiple views.")
prepared = preprocess_images(views)
if prepared is None:
raise gr.Error("Image preprocessing failed. Please try different images.")
strip_ground_slab = any(has_cast_shadow(view) for view in views)
return prepared, partial(_run_original_pipeline_multi, strip_ground_slab=strip_ground_slab), len(prepared)
def generation_gpu_duration(*args, **kwargs):
return GPU_ALLOCATION_SECONDS
def extraction_gpu_duration(*args, **kwargs):
texture_size = kwargs.get("texture_size")
if texture_size is None and len(args) > 2:
texture_size = args[2]
size = int(texture_size or 0)
if size >= MAX_TEXTURE_SIZE:
return GPU_ALLOCATION_SECONDS_MAX_TEXTURE
if size >= HIGH_TEXTURE_SIZE:
return GPU_ALLOCATION_SECONDS_HIGH_TEXTURE
if size > STANDARD_TEXTURE_SIZE:
return GPU_ALLOCATION_SECONDS
return GPU_ALLOCATION_SECONDS_STANDARD_TEXTURE
def gpu_task_duration(mode, *args, **kwargs):
if mode == GPU_BAKE:
return extraction_gpu_duration(*args, **kwargs)
return generation_gpu_duration(*args, **kwargs)
@spaces.GPU(duration=gpu_task_duration, size="large")
def gpu_task(mode, *args, **kwargs):
task = bake_task if mode == GPU_BAKE else preview_task
return task(*args, **kwargs)
def bake_task(state, mesh_simplify, texture_size, progress=None, session="", prebake=False):
with gpu_call(f"gpu_task:{GPU_BAKE}", session, texture_size, **({"prebake": True} if prebake else {})):
if progress is not None:
progress(0.35, desc="Texturing 3D model")
return bake_glb_on_gpu(state, mesh_simplify, texture_size, progress)
def preview_task(*generation_args, req=None, progress=None, sampler=None):
texture_size = generation_args[6]
with gpu_call(f"gpu_task:{GPU_PREVIEW}", session_tag(session_hash(req)), texture_size):
return render_preview_on_gpu(*generation_args, req, progress, sampler)
def texture_model_for_export(
state: Dict[str, Any],
mesh_simplify: float,
texture_size: int,
progress=gr.Progress(),
session: str = "",
prebake: bool = False,
):
return gpu_task(GPU_BAKE, state, mesh_simplify, texture_size, progress, session=session, prebake=prebake)
def generate_preview_frames(
image: Optional[Image.Image],
seed: int,
ss_guidance_strength: float,
ss_sampling_steps: int,
slat_guidance_strength: float,
slat_sampling_steps: int,
texture_size: int = GLB_EXPORT_TEXTURE_SIZE,
req: gr.Request = None,
progress=gr.Progress(),
sampler: Optional[Dict[str, Any]] = None,
):
return gpu_task(
GPU_PREVIEW, image, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength,
slat_sampling_steps, texture_size, req=req, progress=progress, sampler=sampler,
)
def render_preview_on_gpu(
image, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps,
texture_size, req, progress, sampler,
):
if image is None or pipeline is None:
return None, None
if not torch.cuda.is_available():
raise gr.Error("GPU is not ready. Please retry in a few seconds.")
job_start = time.time()
progress(0.02, desc="Preparing images" if isinstance(image, (list, tuple)) else "Preparing image")
torch.cuda.reset_peak_memory_stats()
preprocess_start = time.time()
input_images, original_runner, view_count = _prepare_generation_input(image)
if sampler is not None and original_runner.func is _run_original_pipeline:
original_runner = partial(original_runner, sampler=sampler)
elapsed = time.time() - preprocess_start
logger.info("Preprocess Time: %.2fs", elapsed)
args = (input_images, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps)
try:
return _generate_and_render_model(
args, req, progress, job_start, texture_size=texture_size, original_runner=original_runner,
)
except RuntimeError as error:
if "out of memory" in str(error).lower():
message = ("GPU out of memory. Try fewer images or a simpler object." if view_count > 1
else "GPU out of memory. Please retry with another image.")
raise gr.Error(message) from error
raise gr.Error("Generation failed. Try another image or lower complexity.") from error
except gr.Error:
raise
except Exception as error:
raise gr.Error("Generation failed. Please retry.") from error
finally:
clear_cuda_cache()
if TRELLIS_BENCH:
def bench_bake_gpu_duration(state, configs):
return min(BENCH_MAX_GPU_SECONDS, BENCH_GPU_SECONDS_PER_BAKE * len(configs))
@spaces.GPU(duration=bench_bake_gpu_duration, size="large")
def bench_bakes_on_gpu(state, configs):
from export_profiles import balanced_export_settings, simplify_ratio_for_texture_size
started = time.perf_counter()
original_state = unpack_original_state(state)
torch.cuda.synchronize()
unpack_seconds = round(time.perf_counter() - started, 3)
mesh, app_rep = original_state["mesh"], original_state["gaussian"]
results = []
for config in configs:
overrides = dict(config)
label = overrides.pop("label")
texture_size = int(overrides.pop("texture_size", GLB_EXPORT_TEXTURE_SIZE))
simplify = simplify_ratio_for_texture_size(texture_size, len(mesh.faces))
settings = {**balanced_export_settings(simplify, texture_size), **overrides}
timings = {"unpack": unpack_seconds}
torch.cuda.synchronize()
started = time.perf_counter()
glb = trellis_postprocessing_utils.to_glb(app_rep, mesh, verbose=False, timings=timings, **settings)
torch.cuda.synchronize()
timings["to_glb"] = round(time.perf_counter() - started, 3)
logger.info("BAKE_TIMINGS %s %s", label, json.dumps(timings))
results.append((label, settings, timings, glb))
clear_cuda_cache()
return results
def bench_generate(image, overrides_json, req: gr.Request, progress=gr.Progress(track_tqdm=True)):
overrides = json.loads(overrides_json or "{}")
image, texture_size = validated_generation_input(image, int(overrides.get("texture_size", GLB_EXPORT_TEXTURE_SIZE)))
steps = overrides.get("steps", [12, 12])
frames, state = generate_preview_frames(
image, int(overrides.get("seed", 42)), 6.0, int(steps[0]), 3.0, int(steps[1]), texture_size,
req=req, progress=progress, sampler=overrides.get("sampler"),
)
_, _, fps, _ = get_preview_settings(req)
timings = dict(state.pop("bench_timings", {}))
return write_preview_video(frames, fps=fps), state, json.dumps(timings)
def bench_bakes(state, configs_json):
if not state or state.get("runtime") != "original":
raise gr.Error("Generate with the original runtime first.")
results = bench_bakes_on_gpu(state, json.loads(configs_json))
paths, report = [], []
for label, settings, timings, glb in results:
path = os.path.join(EXPORT_DIR, f"bench-{label}-{uuid4().hex[:8]}.glb")
glb.export(path)
paths.append(path)
report.append({"label": label, "settings": settings, "timings": timings, "file": os.path.basename(path)})
return paths, json.dumps(report)
def validate_texture_size(value):
if value not in GLB_TEXTURE_SIZES:
raise gr.Error("Choose a texture quality of 512, 1024, 1536, 2048, or 4096px.")
return int(value)
def validated_generation_input(image: Any, texture_size: Any):
if isinstance(image, (list, tuple)):
image = validate_multi_images(image)
return image, validate_texture_size(texture_size)
def generation_ready(image: Any, texture_size: Any) -> bool:
if image is None:
return False
try:
validated_generation_input(image, texture_size)
except gr.Error:
return False
return True
def generate_and_extract(
image: Optional[Image.Image],
seed: int,
ss_guidance_strength: float,
ss_sampling_steps: int,
slat_guidance_strength: float,
slat_sampling_steps: int,
texture_size: int = GLB_EXPORT_TEXTURE_SIZE,
req: gr.Request = None,
progress=gr.Progress(),
):
if image is None:
return None, None
image, texture_size = validated_generation_input(image, texture_size)
args = (
image, seed, ss_guidance_strength, ss_sampling_steps,
slat_guidance_strength, slat_sampling_steps, texture_size,
)
access = ACCESS.authorize(req, session_hash(req))
ACCESS.remember(session_hash(req), access.notice)
if access.mode == "denied":
return None, None
started = time.time()
try:
frames, state = generate_preview_frames(*args, req=req, progress=progress)
except Exception as error:
if not is_quota_denial(error):
ACCESS.generation_failed(session_hash(req), access)
raise
notice = ACCESS.quota_denied(access)
if notice:
ACCESS.remember(session_hash(req), notice)
return None, None
try:
defer_or_fail(args, req, access.nonce, error)
except gr.Error:
ACCESS.generation_failed(session_hash(req), access)
raise
return None, None
if not frames:
return None, state
_, _, fps, _ = get_preview_settings(req)
video_path = write_preview_video(frames, fps=fps)
if video_path:
save_run_async(image, video_path, seed=str(seed), source_path="generate")
labs_event("preview_done", session_tag(session_hash(req)), texture_size=texture_size,
took_s=round(time.time() - started, 2))
return video_path, state
def plain_error_text(error: gr.Error) -> str:
text = html.unescape(re.sub(r"<[^>]+>", "", str(error.message)))
return re.sub(r"\s+", " ", text).strip() or "Export failed. Please try again."
def prepare_export_payload(state, mesh_simplify, texture_size, progress=None, session="", prebake=False):
return resolve_export_payload(state, mesh_simplify, texture_size, progress, session, prebake)[0]
def resolve_export_payload(state, mesh_simplify, texture_size, progress=None, session="", prebake=False):
texture_size = validate_texture_size(texture_size)
generation = state.setdefault(GLB_GENERATION_KEY, uuid4().hex)
cached = partial(cached_glb_payload, state, mesh_simplify, texture_size)
joined = False
while True:
payload, flight, owner = FLIGHTS.claim(generation, session, prebake, cached)
if payload is not None:
state.pop(GLB_ERROR_KEY, None)
return payload, joined
if owner:
break
if prebake or flight is None:
return None, joined
joined = True
if progress is not None:
progress(0.35, desc="Texturing 3D model")
flight.done.wait()
payload = None
try:
payload = bake_export_payload(state, mesh_simplify, texture_size, progress, session, prebake)
finally:
FLIGHTS.finish(flight, payload and partial(store_export_payload, state, mesh_simplify, texture_size, payload, prebake))
return payload, joined
def bake_export_payload(state, mesh_simplify, texture_size, progress, session, prebake):
state.pop(GLB_MESH_KEY, None)
glb, failure = texture_model_for_export(state, mesh_simplify, texture_size, progress, session=session, prebake=prebake)
if glb is None:
raise gr.Error(failure or GLB_PREP_FAILURE_MESSAGE)
packaged = cache_prepared_glb({GLB_MESH_KEY: glb}, progress)
payload = packaged.get(GLB_PAYLOAD_KEY)
if payload is None:
raise gr.Error(packaged.get(GLB_ERROR_KEY) or GLB_PACKAGING_FAILURE_MESSAGE)
return payload
def store_export_payload(state, mesh_simplify, texture_size, payload, prebake):
state[GLB_SETTINGS_KEY] = {"simplify": float(mesh_simplify), "texture_size": texture_size}
state[GLB_PAYLOAD_KEY] = payload
state[GLB_PREBAKED_KEY] = prebake
state.pop(GLB_ERROR_KEY, None)
def prebake_glb(state, mesh_simplify, texture_size, req: gr.Request):
if not PREBAKE_GLB or not state or state.get(GENERATED_VIA_KEY) == "pro":
return
run_prebake(state, mesh_simplify, texture_size, session_tag(session_hash(req)),
ip_token=bool(request_header(req, "x-ip-token")))
def run_prebake(state, mesh_simplify, texture_size, session, **fields):
started, payload = time.time(), None
try:
payload = prepare_export_payload(state, mesh_simplify, texture_size, None, session, prebake=True)
except Exception as error:
logger.warning("GLB pre-bake failed: %s", plain_error_text(error) if isinstance(error, gr.Error) else error)
finally:
labs_event("prebake_done", session, ok=payload is not None, texture_size=texture_size,
took_s=round(time.time() - started, 2), **fields)
def start_pro_prebake(entry, state):
if not PREBAKE_GLB or not state:
return
requested = state.get(GLB_REQUESTED_KEY) or {}
args = (
state,
requested.get("simplify", GLB_EXPORT_SIMPLIFY),
requested.get("texture_size", GLB_EXPORT_TEXTURE_SIZE),
session_tag(getattr(entry, "session_hash", "")),
)
context = contextvars.copy_context()
threading.Thread(target=context.run, args=(run_prebake, *args), kwargs={"via": "pro"}, daemon=True).start()
def extract_model(
state: Optional[Dict[str, Any]],
mesh_simplify: float,
texture_size: int,
progress=gr.Progress(),
) -> Tuple[Dict[str, Any], Optional[Dict[str, Any]]]:
if state is None:
return gr.update(), None
started, session, payload, joined = time.time(), current_session_tag(), None, False
try:
payload, joined = resolve_export_payload(state, mesh_simplify, texture_size, progress, session)
except gr.Error as error:
return export_failure_status(state, plain_error_text(error))
except Exception as error:
logger.error("Export allocation failed: %s", error, exc_info=True)
return export_failure_status(state, GLB_PREP_FAILURE_MESSAGE)
finally:
labs_event("glb_done", session, ok=payload is not None, texture_size=texture_size,
took_s=round(time.time() - started, 2), **prebake_source(state, payload, joined))
return gr.update(value="Finalizing GLB for download..."), payload
def prebake_source(state, payload, joined) -> Dict[str, str]:
if payload is None or not state.get(GLB_PREBAKED_KEY):
return {}
return {"prebake": "joined" if joined else "hit"}
def export_failure_status(state: Dict[str, Any], message: str):
state.pop(GLB_MESH_KEY, None)
state[GLB_ERROR_KEY] = message
if is_quota_denial(message):
logger.warning("Export allocation denied; awaiting the signed allowance.")
return gr.update(value=ALLOWANCE_ERROR, visible=True), None
return gr.update(value=safe_message(message, GLB_PREP_FAILURE_MESSAGE), visible=True), None
def defer_pending_fallback(operation: str, args, req: Optional[gr.Request], nonce: Optional[str] = None) -> bool:
if not is_eligible_embed(req):
return False
if FALLBACK.defer(operation, tuple(args), session_hash(req), nonce) is None:
raise gr.Error(CAPACITY_ERROR)
return True
def defer_or_fail(args, req: Optional[gr.Request], nonce: Optional[str], cause: Optional[Exception] = None):
if not defer_pending_fallback(GENERATE, args, req, nonce):
raise gr.Error(ALLOWANCE_ERROR) from cause
def issue_grant_request(image, texture_size, req: gr.Request):
if not generation_ready(image, texture_size) or not grants_required(req):
return None
return ACCESS.grants.request(session_hash(req))
def accept_generation_grant(ack, receipt, req: gr.Request):
ACCESS.grants.settle(session_hash(req), ack, receipt)
def settle_access(req: gr.Request):
notice = ACCESS.pop_notice(session_hash(req))
return gr.update(value=notice or "", visible=bool(notice)), ACCESS.refund_receipt(session_hash(req))
def acknowledge_refund(ack, receipt):
return None
def run_pro_generate(entry, req: gr.Request):
entry.progress_event = current_event_id()
frames, state = generate_preview_frames(*entry.args, req=req)
if state:
state[GENERATED_VIA_KEY] = "pro"
if not frames:
return None, state
start_pro_prebake(entry, state)
_, _, fps, _ = get_preview_settings(req)
video_path = write_preview_video(frames, fps=fps)
if video_path:
save_run_async(entry.args[0], video_path, source_path="pro_fallback")
labs_event("preview_done", session_tag(getattr(entry, "session_hash", "")), texture_size=entry.args[6], via="pro")
return video_path, state
def run_pro_export(entry, req: gr.Request):
entry.progress_event = current_event_id()
state, mesh_simplify, texture_size = entry.args
session = session_tag(getattr(entry, "session_hash", ""))
payload = prepare_export_payload(state, mesh_simplify, texture_size, gr.Progress(), session=session)
labs_event("glb_done", session, ok=True, texture_size=texture_size, via="pro")
return payload
def mirror_fallback_progress(receipt, req: gr.Request, progress=gr.Progress()):
entry = FALLBACK.outputs_for(receipt, session_hash(req)) if receipt else None
if entry is not None:
mirror_progress(entry, progress)
def run_pro_fallback(ticket, permit, req: gr.Request):
return FALLBACK.execute_permit(ticket, permit, req)
def reset_fallback_state(req: gr.Request):
FALLBACK.invalidate_session(session_hash(req))
FLIGHTS.discard_session(session_tag(session_hash(req)))
return None, gr.update(value="", visible=False)
def issue_generation_receipt(req: gr.Request):
if not is_eligible_embed(req):
return None
return FALLBACK.issue(session_hash(req), GENERATE)
def export_quota_denied(state, status) -> bool:
if is_quota_denial(status):
return True
return isinstance(state, dict) and is_quota_denial(state.get(GLB_ERROR_KEY))
def issue_export_receipt(state, mesh_simplify, texture_size, status, req: gr.Request):
if state is None or not export_quota_denied(state, status) or not is_eligible_embed(req):
return None, gr.skip()
entry = FALLBACK.defer(EXPORT, (state, mesh_simplify, texture_size), session_hash(req))
if entry is None:
return None, gr.update(value=CAPACITY_ERROR, visible=True)
return FALLBACK.issue(session_hash(req), EXPORT), gr.update(value=WAITING_STATUS, visible=True)
def fallback_entry(ack, receipt, req: Optional[gr.Request], operation: str):
if not isinstance(receipt, dict) or not receipt.get("ticket"):
return None, None
if not isinstance(ack, dict) or ack.get("ticket") != receipt.get("ticket"):
return None, None
entry = FALLBACK.outputs_for(receipt, session_hash(req))
if entry is None or entry.operation != operation or entry.cancelled:
return None, None
if ack.get("cancelled") is True or ack.get("ready") is not True:
return None, safe_message(ack.get("error")) if ack.get("error") else None
return entry, None
def restore_generation_outputs(ack, receipt, req: gr.Request):
entry, error = fallback_entry(ack, receipt, req, GENERATE)
if entry is not None and entry.status == "complete" and entry.outputs is not None:
video_path, state = entry.outputs
return video_path, state, gr.update(value="", visible=False)
if entry is not None or error:
return gr.skip(), gr.skip(), gr.update(value=error or ALLOWANCE_ERROR, visible=True)
return gr.skip(), gr.skip(), gr.skip()
def restore_export_outputs(ack, receipt, req: gr.Request):
entry, error = fallback_entry(ack, receipt, req, EXPORT)
if entry is not None and entry.status == "complete" and entry.outputs is not None:
return gr.update(value="", visible=False), entry.outputs, gr.update(value="", visible=False)
if entry is not None or error:
return gr.update(value=error or ALLOWANCE_ERROR, visible=True), gr.skip(), gr.update(value="", visible=False)
return gr.skip(), gr.skip(), gr.skip()
FALLBACK.register(GENERATE, run_pro_generate)
FALLBACK.register(EXPORT, run_pro_export)
ACCESS = AccessPolicy(GrantBook(), public_allowance_from_env())
fallback_bridge = bridge_script(Path(__file__).with_name("fallback_bridge.js"))
grant_bridge = bridge_script(Path(__file__).with_name("grant_bridge.js"))
refund_bridge = bridge_script(Path(__file__).with_name("refund_bridge.js"))
progress_bar_js = Path(__file__).with_name("progress_bar.js").read_text()
GLB_PROGRESS_SECONDS = {512: 4, 1024: 5, 1536: 8, 2048: 13, 4096: 40}
def progress_bar(op: str, seconds: float, done: str, stages, source: str = "", gpu: bool = False):
return gr.HTML(
value=(
f''
),
js_on_load=progress_bar_js,
elem_id=f"{op}-progress",
elem_classes="labs-progress-host",
)
def progress_hook(call: str) -> str:
return f"(...args) => {{ const bars = window.labsProgressBars || {{}}; try {{ {call} }} catch (e) {{}} return args; }}"
def finish_progress(*_):
return None
css = """
.labs-access-notice {
margin-top: 8px;
padding: 12px 16px;
border: 1px solid #5A11FF;
border-radius: 8px;
background: #f4f0ff;
color: #1a1a1a;
font-size: 14px;
line-height: 1.5;
}
.labs-access-notice a {
display: inline-block;
margin-left: 4px;
padding: 4px 12px;
border-radius: 10px;
background: #5A11FF;
color: #fff !important;
font-weight: 600;
text-decoration: none;
}
/* Force light page background even when the viewer's OS is in dark mode */
body { background: #f9fafb !important; }
h1, h2, h3 { text-align: center; display: block; }
h1 a {color: #5A11FF !important; text-decoration: none !important;}
footer { visibility: hidden; }
.gradio-container { width: 100%; max-width: 1100px !important; margin: 0 auto; }
@media (max-width: 767px) {
.generation-input-row { flex-direction: column !important; }
.generation-input-row > div { width: 100% !important; flex: none !important; }
}
.gr-image-container { display: flex !important; justify-content: center !important; align-items: center !important; width: 100%; height: 240px; }
.gr-image-container img { width: 100%; height: 100%; object-fit: contain; object-position: center; }
.gr-image { display: flex; justify-content: center; align-items: center; width: 100%; height: 512px; overflow: hidden; }
.gr-image img { width: 100%; height: 100%; object-fit: cover; object-position: center; }
.video-container video { width: 100% !important; height: 100% !important; object-fit: contain !important; object-position: center !important; }
.sponsor-banner { text-align: center; margin: 8px 0 14px 0; }
.sponsor-banner-title { font-size: 1.05rem; font-weight: 700; margin-bottom: 8px; }
.sponsor-banner-title a {
color: #5A11FF !important;
text-decoration: none !important;
}
.sponsor-banner-button {
display: inline-block;
padding: 8px 14px;
border-radius: 10px;
font-weight: 700;
text-decoration: none !important;
background: linear-gradient(90deg, #2f6bff 0%, #7d4dff 100%);
color: #ffffff !important;
}
.inter-app-nav {
text-align: center;
margin-bottom: 15px;
border-bottom: 1px solid #eee;
padding-bottom: 10px;
}
.inter-app-nav-link {
color: #2f6bff !important;
font-weight: bold !important;
margin: 0 12px !important;
text-decoration: none !important;
}
.toast-wrap, .toast-body, .toast-container { display: none !important; }
.model-container .progress-text, .model-container .progress-level { display: none !important; }
.video-container .progress-text, .video-container .progress-level, .video-container .eta-bar { display: none !important; }
.labs-progress-host:has(.labs-progress[hidden]) { display: none !important; }
.labs-progress-host { opacity: 1 !important; max-width: 680px; width: 100%; margin: 0 auto !important; }
.labs-progress-host [data-testid="status-tracker"], .labs-progress-host .wrap { display: none !important; }
.labs-progress { padding: 4px 0; text-align: center; }
.labs-progress-label { color: #475569; font-size: 14px; font-weight: 600; margin-bottom: 6px; font-variant-numeric: tabular-nums; }
.labs-progress-track { height: 8px; border-radius: 999px; background: #e2e2e2; overflow: hidden; }
.labs-progress-fill { width: 0; height: 100%; border-radius: 999px; background: linear-gradient(90deg, #2f6bff 0%, #7d4dff 100%); transition: width 0.2s linear; }
@media screen and (min-width: 768px) {
.gr-image-container { height: 360px !important; }
.video-container { height: 360px !important; max-width: 680px !important; margin: 0 auto !important; aspect-ratio: auto !important; }
.model-container { height: 480px !important; max-width: 680px !important; margin: 0 auto !important; }
}
.custom-header { display: flex; align-items: center; height: 100%; }
.stl-download-btn { max-width: 680px !important; margin: 0 auto !important; }
.stl-download-btn button { width: 100% !important; background: #5A11FF !important; border-color: #5A11FF !important; color: #fff !important; }
"""
schema_data = {
"@context": "https://schema.org",
"@type": "SoftwareApplication",
"name": "Pocket 3D AI by SkyeBrowse",
"operatingSystem": "Web",
"applicationCategory": "MultimediaApplication",
"description": "Instant AI-powered 3D model generation from a single image. Upload a photo and get a downloadable GLB and STL model in seconds.",
"author": {
"@type": "Organization",
"name": "SkyeBrowse",
"url": "https://www.skyebrowse.com"
},
"offers": {
"@type": "Offer",
"price": "0",
"priceCurrency": "USD"
}
}
PAGE_TITLE = "Pocket 3D AI: Free Image to 3D Model Generator (GLB, STL) | SkyeBrowse"
PAGE_DESCRIPTION = "Turn one photo or 2-6 views into a textured 3D model in your browser. Free, no install. Download GLB for Blender and game engines or STL for 3D printing."
PAGE_IMAGE = "https://skyebrowse-prod-web-assets.s3.us-east-2.amazonaws.com/pocket_3d_ai_hero_adac81efac.jpg"
head_html = f"""
"""
custom_js = """
() => {
// Force light mode. Gradio applies dark via a `dark` class; in
// web-component embeds (e.g. app.skyebrowse.com) that class lands on the host
// element, not /, so strip it from EVERY element that has it.
function forceLight() {
document.querySelectorAll('.dark').forEach(el => el.classList.remove('dark'));
document.documentElement.classList.remove('dark');
}
forceLight();
new MutationObserver(forceLight).observe(document.documentElement, {
attributes: true, attributeFilter: ['class'], subtree: true
});
new MutationObserver(() => {
document.querySelectorAll('.progress-text, .eta-bar, .progress-level-inner').forEach(el => {
if (el.textContent.match(/zero\\s*gpu/i)) {
el.style.visibility = 'hidden';
}
});
}).observe(document.body, {childList: true, subtree: true, characterData: true});
// Rewrite external app links when hosted on *.app.skyebrowse.com
const hostname = window.location.hostname;
if (hostname.endsWith('app.skyebrowse.com')) {
const origin = window.location.origin;
const linkMap = {
'interiorai.skyebrowse.com': origin + '/interior-ai',
'anime.skyebrowse.com': origin + '/anime-ai',
'3dai.skyebrowse.com': origin + '/3d-ai',
'crowd.skyebrowse.com': origin + '/crowd-counter',
'app.skyebrowse.com': origin + '/app',
'www.skyebrowse.com': origin,
};
function rewriteLinks() {
document.querySelectorAll('a[href]').forEach(a => {
try {
const originalHref = a.dataset.appOriginalHref || a.href;
const url = new URL(originalHref);
if (linkMap[url.hostname]) {
a.dataset.appOriginalHref = originalHref;
a.href = linkMap[url.hostname];
if (a.classList.contains('inter-app-nav-link')) {
a.target = '_self';
}
}
} catch(e) {}
});
}
rewriteLinks();
new MutationObserver(rewriteLinks).observe(document.body, {childList: true, subtree: true});
}
}
"""
with gr.Blocks(title=PAGE_TITLE) as demo:
default_ss_steps = 12
default_slat_steps = 12
# with gr.Row(equal_height=True):
# gr.Image("assets/sb_pocket_logo_dark.png", show_label=False, container=False, show_download_button=False, min_width=50, interactive=False, show_fullscreen_button=False)
gr.Markdown("# 🤖 Pocket 3D AI by [SkyeBrowse](https://www.skyebrowse.com)")
gr.Markdown(
"Turn any image into a 3D model in seconds. Upload or paste a photo and Pocket 3D AI will generate a "
"full 3D model with textures. Preview it right in your browser, then download the **GLB** or **STL** file "
"for 3D printing, game engines, AR/VR, and more."
)
gr.HTML(
''
)
with gr.Row():
gr.HTML(
''
)
with gr.Column():
with gr.Row(elem_classes="generation-input-row"):
with gr.Column(scale=2, min_width=100, variant="default"):
texture_size = gr.Dropdown(
choices=[(f"{size}px", size) for size in GLB_TEXTURE_SIZES],
value=GLB_EXPORT_TEXTURE_SIZE, label="Texture quality",
info="Choose before uploading, or regenerate to apply a change.",
)
with gr.Tabs():
with gr.Tab("Single Image", id=0) as single_image_tab:
image_prompt = gr.Image(
label="Single Image Input",
format="png",
image_mode="RGBA",
type="pil",
sources=["upload", "clipboard"],
container=True,
webcam_options=gr.WebcamOptions(mirror=True),
visible=True,
height=240,
elem_classes="gr-image-container",
)
regenerate_button = gr.Button("Regenerate", variant="secondary")
with gr.Tab(
"Multiple Images", id=1, visible=TRELLIS_RUNTIME == "original",
) as multi_image_tab:
multi_image_gallery = gr.Gallery(
label="Multiple Images (2–6 views)",
format="png",
type="pil",
columns=3,
rows=2,
height=240,
object_fit="contain",
buttons=[],
)
multi_image_button = gr.Button(
"Generate from Multiple Images", variant="primary", size="lg",
)
gr.Markdown(
"Add 2–6 views of the same object. Extra angles resolve shape "
"that a single photo leaves ambiguous."
)
with gr.Column(scale=5, min_width=100):
video_output = gr.Video(
label="Preview",
height=240,
elem_classes="video-container",
visible=True,
autoplay=True,
loop=True,
buttons=["download"],
interactive=False,
)
preview_progress = progress_bar(
"preview", 5, "Preview ready",
[[0, "Preparing image"], [15, "Generating 3D structure"], [60, "Rendering preview"]],
source=".video-container .progress-level-inner", gpu=True,
)
access_notice = gr.HTML(value="", visible=False, elem_id="labs-access-notice")
with gr.Row(equal_height=False):
with gr.Column(scale=2, min_width=100, variant="default"):
with gr.Group(visible=True) as single_examples_group:
examples = gr.Examples(
examples=[f"./assets/example_image/{image}" for image in os.listdir("./assets/example_image")],
inputs=[image_prompt],
examples_per_page=9,
)
with gr.Group(visible=False) as multi_examples_group:
multi_example_images = prepare_multi_example()
if multi_example_images:
gr.Examples(
examples=multi_example_images,
inputs=[image_prompt],
fn=split_multi_example,
outputs=[multi_image_gallery],
run_on_click=True,
cache_examples=False,
examples_per_page=8,
)
with gr.Column(scale=5):
model_output = LitModel3D(
label="",
container=True,
zoom_speed=0.5,
pan_speed=3.0,
exposure=10.0,
height=480,
elem_classes="model-container",
visible=False,
)
glb_progress = progress_bar(
"glb", GLB_PROGRESS_SECONDS[GLB_EXPORT_TEXTURE_SIZE], "3D model ready",
[[0, "Starting 3D model"], [10, "Texturing 3D model"], [80, "Packaging GLB"]],
source="#glb-progress .progress-level-inner", gpu=True,
)
glb_button = gr.Button(
"Generate 3D Model", interactive=False, size="lg", variant="primary",
elem_classes="stl-download-btn",
)
glb_download_button = gr.DownloadButton(
"Download 3D Model", visible=False, size="lg", variant="primary",
elem_classes="stl-download-btn", elem_id="glb-download-btn",
)
glb_export_status = gr.Textbox(
label="GLB export", interactive=False, visible=False, elem_id="glb-export-status",
)
stl_download_button = gr.DownloadButton(
label="Generate STL",
visible=True,
interactive=False,
size="lg",
variant="primary",
elem_classes="stl-download-btn",
elem_id="stl-download-btn",
)
progress_bar("stl", 3, "STL ready", [[0, "Building STL"]])
with gr.Accordion(label="Generation Settings", open=False, visible=show_options and not RUNNING_ON_SPACES):
seed_slider = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
gr.Markdown("Stage 1: Sparse Structure Generation")
with gr.Row():
ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=default_ss_steps, step=1)
gr.Markdown("Stage 2: Structured Latent Generation")
with gr.Row():
slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=default_slat_steps, step=1)
if RUNNING_ON_SPACES:
seed_slider = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1, visible=False)
randomize_seed = gr.Checkbox(label="Randomize Seed", value=True, visible=False)
ss_guidance_strength = gr.Slider(0.0, 10.0, value=6.0, step=0.1, visible=False)
ss_sampling_steps = gr.Slider(1, 50, value=default_ss_steps, step=1, visible=False)
slat_guidance_strength = gr.Slider(0.0, 10.0, value=3.0, step=0.1, visible=False)
slat_sampling_steps = gr.Slider(1, 50, value=default_slat_steps, step=1, visible=False)
mesh_simplify = gr.Number(value=GLB_EXPORT_SIMPLIFY, visible=False)
export_texture_size = gr.Number(value=GLB_EXPORT_TEXTURE_SIZE, precision=0, visible=False)
export_payload_state = gr.State(value=None)
materialized_glb_state = gr.State(value=None)
stl_path_state = gr.State(value=None)
demo.load(start_session)
demo.unload(end_session)
pipeline_state = gr.State(value=None)
fallback_receipt = gr.JSON(visible=False)
fallback_ack = gr.JSON(visible=False)
fallback_status = gr.Markdown(value="", visible=False, elem_id="labs-fallback-status")
grant_receipt = gr.JSON(visible=False)
grant_ack = gr.JSON(visible=False)
refund_receipt = gr.JSON(visible=False)
refund_ack = gr.JSON(visible=False)
fallback_ticket_api = gr.Textbox(visible=False)
fallback_permit_api = gr.Textbox(visible=False)
fallback_api_trigger = gr.Button(visible=False)
fallback_api_trigger.click(
fn=run_pro_fallback,
inputs=[fallback_ticket_api, fallback_permit_api],
outputs=[fallback_ack],
api_name="run_pro_fallback",
show_progress="hidden",
)
combined_inputs = [
image_prompt,
seed_slider,
ss_guidance_strength,
ss_sampling_steps,
slat_guidance_strength,
slat_sampling_steps,
texture_size,
]
multi_inputs = [
multi_image_gallery,
seed_slider,
ss_guidance_strength,
ss_sampling_steps,
slat_guidance_strength,
slat_sampling_steps,
texture_size,
]
extract_inputs = [
pipeline_state,
mesh_simplify,
export_texture_size,
]
reset_outputs = [
video_output, pipeline_state, model_output, export_payload_state,
materialized_glb_state, stl_path_state, stl_download_button, glb_button,
glb_export_status, glb_download_button,
]
def wire_generation_flow(trigger, generation_inputs):
receipt = (
trigger(
begin_generation,
inputs=[randomize_seed, seed_slider],
outputs=[seed_slider, *reset_outputs],
show_progress="hidden",
trigger_mode="always_last",
js=progress_hook("Object.values(bars).forEach(bar => bar.cancel()); bars.preview?.start();"),
)
.then(reset_fallback_state, outputs=[fallback_receipt, fallback_status], show_progress="hidden")
.then(
fn=issue_grant_request,
inputs=[generation_inputs[0], generation_inputs[6]],
outputs=[grant_receipt],
show_progress="hidden",
)
.then(
fn=accept_generation_grant,
inputs=[grant_ack, grant_receipt],
js=grant_bridge,
show_progress="minimal",
)
.then(
fn=generate_and_extract,
inputs=generation_inputs,
outputs=[video_output, pipeline_state],
show_progress="full",
scroll_to_output=True,
)
.then(fn=settle_access, outputs=[access_notice, refund_receipt], show_progress="hidden")
.then(
fn=acknowledge_refund,
inputs=[refund_ack, refund_receipt],
js=refund_bridge,
show_progress="hidden",
)
.then(fn=issue_generation_receipt, outputs=[fallback_receipt], show_progress="hidden")
)
receipt.then(
fn=mirror_fallback_progress,
inputs=[fallback_receipt],
show_progress="full",
show_progress_on=[video_output],
concurrency_limit=None,
api_visibility="private",
)
return (
receipt.then(
fn=restore_generation_outputs,
inputs=[fallback_ack, fallback_receipt],
outputs=[video_output, pipeline_state, fallback_status],
js=fallback_bridge,
show_progress="minimal",
)
.then(
fn=enable_model_exports,
inputs=[pipeline_state],
outputs=[stl_download_button, glb_button, stl_path_state, export_texture_size],
show_progress="hidden",
js=progress_hook(
"setTimeout(() => bars.preview?.finish(Boolean(document.querySelector('.video-container video'))), 300);"
),
)
.then(
fn=prebake_glb,
inputs=[pipeline_state, mesh_simplify, export_texture_size],
show_progress="hidden",
api_visibility="private",
)
)
for trigger in [partial(gr.on, [image_prompt.upload, regenerate_button.click]), examples.load_input_event.then]:
wire_generation_flow(trigger, combined_inputs)
wire_generation_flow(multi_image_button.click, multi_inputs)
single_image_tab.select(
lambda: (gr.update(visible=True), gr.update(visible=False)),
outputs=[single_examples_group, multi_examples_group],
show_progress="hidden",
)
multi_image_tab.select(
lambda: (gr.update(visible=False), gr.update(visible=True)),
outputs=[single_examples_group, multi_examples_group],
show_progress="hidden",
)
image_prompt.clear(
reset_generation_outputs, outputs=reset_outputs, show_progress="hidden",
js=progress_hook("Object.values(bars).forEach(bar => bar.cancel());"),
).then(
fn=reset_fallback_state,
outputs=[fallback_receipt, fallback_status],
show_progress="hidden",
)
export_receipt = glb_button.click(
begin_glb_export,
outputs=[glb_button, glb_export_status], show_progress="hidden",
).then(
fn=extract_model,
inputs=extract_inputs,
outputs=[glb_export_status, export_payload_state],
show_progress="full",
show_progress_on=[glb_progress],
scroll_to_output=True,
js=progress_hook(f"bars.glb?.start({{seconds: {json.dumps(GLB_PROGRESS_SECONDS)}[args[2]]}});"),
).then(
fn=issue_export_receipt,
inputs=[pipeline_state, mesh_simplify, export_texture_size, glb_export_status],
outputs=[fallback_receipt, glb_export_status],
show_progress="hidden",
)
export_receipt.then(
fn=mirror_fallback_progress,
inputs=[fallback_receipt],
show_progress="full",
show_progress_on=[glb_progress],
concurrency_limit=None,
api_visibility="private",
)
export_receipt.then(
fn=restore_export_outputs,
inputs=[fallback_ack, fallback_receipt],
outputs=[glb_export_status, export_payload_state, fallback_status],
js=fallback_bridge,
show_progress="minimal",
).then(
fn=materialize_export_payload,
inputs=[export_payload_state],
outputs=[model_output, materialized_glb_state],
show_progress="full",
js=progress_hook("bars.glb?.stage('Loading 3D viewer', 90);"),
).then(
fn=finish_glb_export,
inputs=[materialized_glb_state, pipeline_state, glb_export_status],
outputs=[glb_button, glb_export_status, glb_download_button],
show_progress="hidden",
js=progress_hook(
"setTimeout(() => bars.glb?.finish(document.querySelector('.model-container')?.getClientRects().length > 0), 300);"
),
)
stl_download_button.click(
fn=generate_stl_download,
inputs=[pipeline_state, stl_path_state],
outputs=[stl_download_button, stl_path_state],
show_progress="hidden",
js=progress_hook(
"if (!/download/i.test(document.querySelector('#stl-download-btn')?.textContent)) bars.stl?.start();"
),
).then(
fn=finish_progress,
show_progress="hidden",
api_visibility="private",
js=progress_hook(
"setTimeout(() => bars.stl?.finish(/download/i.test(document.querySelector('#stl-download-btn')?.textContent)), 300);"
),
)
if TRELLIS_BENCH:
bench_request = gr.Textbox(visible=False)
bench_report = gr.Textbox(visible=False)
bench_files = gr.File(file_count="multiple", visible=False)
gr.Button(visible=False).click(
bench_generate, inputs=[image_prompt, bench_request],
outputs=[video_output, pipeline_state, bench_report], api_name="bench_generate",
)
gr.Button(visible=False).click(
bench_bakes, inputs=[pipeline_state, bench_request],
outputs=[bench_files, bench_report], api_name="bench_bakes",
)
gr.Markdown(
"""
---
### About Pocket 3D AI by SkyeBrowse
Pocket 3D AI generates 3D models from a single image in seconds.
Upload any photo and get a downloadable **GLB** and **STL** file ready for 3D printing,
game engines, or AR/VR applications. Built on **TRELLIS** image-to-3D reconstruction,
this tool is part of the **SkyeBrowse AI** ecosystem powering
**3D reconstruction**, **photogrammetry**, and **spatial data visualization**.
"""
)
if __name__ == "__main__":
if pipeline is None:
logger.critical("Pipeline failed to initialize. Exiting.")
sys.exit(1)
logger.info("Launching runtime: %s", TRELLIS_RUNTIME)
app_kwargs = gradio_runtime.build_launch_kwargs(
theme=gr.themes.Soft(), css=css, js=custom_js, head=head_html, title=PAGE_TITLE
)
if RUNNING_ON_SPACES:
logger.info("Launching on HuggingFace Spaces")
launch_kwargs = {
**app_kwargs,
"share": True,
"allowed_paths": ALLOWED_PATHS,
}
if PUBLIC_BASE_URL:
launch_kwargs["root_path"] = PUBLIC_BASE_URL
logger.info("Using PUBLIC_BASE_URL for URL generation: %s", PUBLIC_BASE_URL)
else:
logger.warning(
"PUBLIC_BASE_URL is not set. If using a custom domain, set PUBLIC_BASE_URL=https://your-domain"
)
demo.queue(max_size=20, default_concurrency_limit=20, api_open=False).launch(**launch_kwargs)
elif prod:
logger.info("Launching in PRODUCTION mode on port %s", port)
demo.queue(max_size=20, default_concurrency_limit=5).launch(
**app_kwargs,
server_name="0.0.0.0",
server_port=port,
favicon_path="assets/sb_3d_ai_logo.png",
share=False,
allowed_paths=ALLOWED_PATHS,
)
else:
logger.info("Launching in DEVELOPMENT mode on port %s", port)
demo.queue(api_open=False).launch(
**app_kwargs,
server_name="0.0.0.0",
server_port=port,
favicon_path="assets/sb_3d_ai_logo.png",
debug=True,
share=True,
allowed_paths=ALLOWED_PATHS,
)