wehub-resource-sync
559 lines
22 KiB
Python
559 lines
22 KiB
Python
import os
|
|
import subprocess
|
|
import argparse
|
|
import math
|
|
import time
|
|
import shutil
|
|
import cv2
|
|
import torch
|
|
import numpy as np
|
|
import base64
|
|
import io
|
|
import json
|
|
from datetime import datetime
|
|
from typing import *
|
|
from PIL import Image
|
|
|
|
import threading
|
|
try:
|
|
import nest_asyncio
|
|
nest_asyncio.apply()
|
|
except ImportError:
|
|
pass
|
|
|
|
# Lock for model initialization
|
|
init_lock = threading.Lock()
|
|
|
|
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
|
|
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
|
os.environ.setdefault("ATTN_BACKEND", "flash_attn")
|
|
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'
|
|
|
|
import spaces
|
|
from gradio import Server
|
|
from gradio.data_classes import FileData
|
|
from fastapi.responses import HTMLResponse
|
|
from fastapi.staticfiles import StaticFiles
|
|
|
|
from pixal3d.modules.sparse import SparseTensor
|
|
from pixal3d.pipelines import Pixal3DImageTo3DPipeline
|
|
from pixal3d.renderers import EnvMap
|
|
from pixal3d.utils import render_utils
|
|
import o_voxel
|
|
|
|
# ============================================================================
|
|
# Constants & Defaults
|
|
# ============================================================================
|
|
|
|
MAX_SEED = np.iinfo(np.int32).max
|
|
TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
|
|
os.makedirs(TMP_DIR, exist_ok=True)
|
|
|
|
MODES = [
|
|
{"name": "Normal", "icon": "assets/app/normal.png", "render_key": "normal"},
|
|
{"name": "Clay render", "icon": "assets/app/clay.png", "render_key": "clay"},
|
|
{"name": "Base color", "icon": "assets/app/basecolor.png", "render_key": "base_color"},
|
|
{"name": "HDRI forest", "icon": "assets/app/hdri_forest.png", "render_key": "shaded_forest"},
|
|
{"name": "HDRI sunset", "icon": "assets/app/hdri_sunset.png", "render_key": "shaded_sunset"},
|
|
{"name": "HDRI courtyard", "icon": "assets/app/hdri_courtyard.png", "render_key": "shaded_courtyard"},
|
|
]
|
|
STEPS = 8
|
|
|
|
# Cascade parameters
|
|
CASCADE_LR_RESOLUTION = 512
|
|
CASCADE_MAX_NUM_TOKENS = 49152
|
|
|
|
# MoGe defaults
|
|
MOGE_MODEL_NAME = "Ruicheng/moge-2-vitl"
|
|
WILD_MESH_SCALE = 1.0
|
|
WILD_EXTEND_PIXEL = 0
|
|
WILD_IMAGE_RESOLUTION = 512
|
|
|
|
# Image Cond Model configs
|
|
IMAGE_COND_CONFIGS = {
|
|
"ss": {
|
|
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
|
"image_size": 512,
|
|
"grid_resolution": 16,
|
|
},
|
|
"shape_512": {
|
|
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
|
"image_size": 512,
|
|
"grid_resolution": 32,
|
|
"use_naf_upsample": True,
|
|
"naf_target_size": 512,
|
|
},
|
|
"shape_1024": {
|
|
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
|
"image_size": 1024,
|
|
"grid_resolution": 64,
|
|
"use_naf_upsample": True,
|
|
"naf_target_size": 512,
|
|
},
|
|
"tex_1024": {
|
|
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
|
"image_size": 1024,
|
|
"grid_resolution": 64,
|
|
"use_naf_upsample": True,
|
|
"naf_target_size": 1024,
|
|
},
|
|
}
|
|
|
|
# ============================================================================
|
|
# Model Loading
|
|
# ============================================================================
|
|
|
|
def build_image_cond_model(config: dict):
|
|
from pixal3d.trainers.flow_matching.mixins.image_conditioned_proj import DinoV3ProjFeatureExtractor
|
|
model = DinoV3ProjFeatureExtractor(**config)
|
|
model.eval()
|
|
return model
|
|
|
|
def load_moge_model(device="cuda", model_name=MOGE_MODEL_NAME):
|
|
from moge.model.v2 import MoGeModel
|
|
moge_model = MoGeModel.from_pretrained(model_name).to(device)
|
|
moge_model.eval()
|
|
return moge_model
|
|
|
|
# Global instances (lazy loaded or loaded at start)
|
|
pipeline = None
|
|
moge_model = None
|
|
envmap = None
|
|
LOW_VRAM = os.environ.get("LOW_VRAM", "0") == "1"
|
|
|
|
def init_models():
|
|
global pipeline, moge_model, envmap
|
|
with init_lock:
|
|
if pipeline is not None:
|
|
return
|
|
|
|
# GPU / CUDA Diagnostics (runs when GPU is allocated)
|
|
import subprocess as _sp
|
|
print("=" * 60)
|
|
print("[Diagnostics] PyTorch version:", torch.__version__)
|
|
print("[Diagnostics] CUDA available:", torch.cuda.is_available())
|
|
if torch.cuda.is_available():
|
|
print("[Diagnostics] CUDA version:", torch.version.cuda)
|
|
print("[Diagnostics] cuDNN version:", torch.backends.cudnn.version())
|
|
for i in range(torch.cuda.device_count()):
|
|
name = torch.cuda.get_device_name(i)
|
|
cap = torch.cuda.get_device_capability(i)
|
|
mem = torch.cuda.get_device_properties(i).total_memory / 1024**3
|
|
print(f"[Diagnostics] GPU {i}: {name}, sm_{cap[0]}{cap[1]}, {mem:.1f} GB")
|
|
try:
|
|
res = _sp.run(["nvidia-smi", "--query-gpu=name,compute_cap,memory.total", "--format=csv,noheader"], capture_output=True, text=True, timeout=10)
|
|
print("[Diagnostics] nvidia-smi:", res.stdout.strip())
|
|
except Exception as e:
|
|
print(f"[Diagnostics] nvidia-smi failed: {e}")
|
|
print("=" * 60)
|
|
|
|
model_path = "TencentARC/Pixal3D"
|
|
print(f"[Pipeline] Loading from {model_path}...")
|
|
pipeline = Pixal3DImageTo3DPipeline.from_pretrained(model_path)
|
|
|
|
print("[ImageCond] Building DinoV3ProjFeatureExtractor models...")
|
|
pipeline.image_cond_model_ss = build_image_cond_model(IMAGE_COND_CONFIGS["ss"])
|
|
pipeline.image_cond_model_shape_512 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_512"])
|
|
pipeline.image_cond_model_shape_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_1024"])
|
|
pipeline.image_cond_model_tex_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["tex_1024"])
|
|
|
|
if LOW_VRAM:
|
|
# Low-VRAM mode: models stay on CPU, loaded to GPU on-demand per stage.
|
|
print("[NAF] Pre-downloading NAF upsampler weights (CPU only)...")
|
|
for attr in ['image_cond_model_ss', 'image_cond_model_shape_512',
|
|
'image_cond_model_shape_1024', 'image_cond_model_tex_1024']:
|
|
m = getattr(pipeline, attr, None)
|
|
if m is not None and getattr(m, 'use_naf_upsample', False):
|
|
m._load_naf()
|
|
pipeline._device = torch.device("cuda")
|
|
pipeline.low_vram = True
|
|
print("[Pipeline] Low-VRAM mode enabled.")
|
|
else:
|
|
# Standard mode: all models loaded to GPU at once.
|
|
pipeline.low_vram = False
|
|
pipeline.cuda()
|
|
pipeline.image_cond_model_ss.cuda()
|
|
pipeline.image_cond_model_shape_512.cuda()
|
|
pipeline.image_cond_model_shape_1024.cuda()
|
|
pipeline.image_cond_model_tex_1024.cuda()
|
|
print("[NAF] Pre-loading NAF upsampler model...")
|
|
for attr in ['image_cond_model_ss', 'image_cond_model_shape_512',
|
|
'image_cond_model_shape_1024', 'image_cond_model_tex_1024']:
|
|
m = getattr(pipeline, attr, None)
|
|
if m is not None and getattr(m, 'use_naf_upsample', False):
|
|
m._load_naf()
|
|
|
|
print("[MoGe-2] Loading model for camera estimation...")
|
|
if LOW_VRAM:
|
|
# Low-VRAM: load MoGe to CPU, move to GPU on-demand per request.
|
|
moge_model = load_moge_model(device="cpu")
|
|
print("[MoGe-2] Low-VRAM mode: MoGe stays on CPU, loaded to GPU on-demand.")
|
|
else:
|
|
moge_model = load_moge_model(device="cuda")
|
|
|
|
print("[EnvMap] Loading environment maps...")
|
|
_base = os.path.dirname(os.path.abspath(__file__))
|
|
_envmap_device = 'cpu' if LOW_VRAM else 'cuda'
|
|
envmap = {
|
|
'forest': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread(os.path.join(_base, 'assets/hdri/forest.exr'), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device=_envmap_device)),
|
|
'sunset': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread(os.path.join(_base, 'assets/hdri/sunset.exr'), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device=_envmap_device)),
|
|
'courtyard': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread(os.path.join(_base, 'assets/hdri/courtyard.exr'), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device=_envmap_device)),
|
|
}
|
|
|
|
# ============================================================================
|
|
# Utilities
|
|
# ============================================================================
|
|
|
|
def compute_f_pixels(camera_angle_x: float, resolution: int) -> float:
|
|
focal_length = 16.0 / torch.tan(torch.tensor(camera_angle_x / 2.0))
|
|
f_pixels = focal_length * resolution / 32.0
|
|
return float(f_pixels.item())
|
|
|
|
def distance_from_fov(camera_angle_x, grid_point, target_point, mesh_scale, image_resolution):
|
|
rotation_matrix = torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]])
|
|
gp = grid_point.to(torch.float32) @ rotation_matrix.T
|
|
gp = gp / mesh_scale / 2
|
|
xw, yw, zw = gp[0].item(), gp[1].item(), gp[2].item()
|
|
xt, yt = float(target_point[0].item()), float(target_point[1].item())
|
|
f_pixels = compute_f_pixels(camera_angle_x, image_resolution)
|
|
x_ndc = xt - image_resolution / 2.0
|
|
y_ndc = -(yt - image_resolution / 2.0)
|
|
distance_x = f_pixels * xw / x_ndc - yw
|
|
return {"distance_from_x": float(distance_x), "f_pixels": float(f_pixels)}
|
|
|
|
def get_camera_params_wild_moge(image_path, device="cuda", mesh_scale=1.0, extend_pixel=0, image_resolution=512):
|
|
pil_image = Image.open(image_path).convert("RGB")
|
|
width, height = pil_image.size
|
|
image_np = np.array(pil_image).astype(np.float32) / 255.0
|
|
image_tensor = torch.from_numpy(image_np).permute(2, 0, 1).to(device)
|
|
if LOW_VRAM:
|
|
moge_model.to(device)
|
|
with torch.no_grad():
|
|
output = moge_model.infer(image_tensor)
|
|
if LOW_VRAM:
|
|
moge_model.cpu()
|
|
torch.cuda.empty_cache()
|
|
intrinsics = output["intrinsics"].squeeze().cpu().numpy()
|
|
fx_normalized = intrinsics[0, 0]
|
|
fx = fx_normalized * width
|
|
camera_angle_x = 2 * math.atan(width / (2 * fx))
|
|
|
|
grid_point = torch.tensor([-1.0, 0.0, 0.0])
|
|
distance = distance_from_fov(
|
|
camera_angle_x, grid_point,
|
|
torch.tensor([0 - extend_pixel, image_resolution - 1 + extend_pixel]),
|
|
mesh_scale, image_resolution
|
|
)["distance_from_x"]
|
|
return {'camera_angle_x': camera_angle_x, 'distance': distance, 'mesh_scale': mesh_scale}
|
|
|
|
def pack_state(shape_slat, tex_slat, res):
|
|
state_data = {
|
|
'shape_slat_feats': shape_slat.feats.cpu().numpy(),
|
|
'tex_slat_feats': tex_slat.feats.cpu().numpy(),
|
|
'coords': shape_slat.coords.cpu().numpy(),
|
|
'res': res,
|
|
}
|
|
import random
|
|
state_path = os.path.join(TMP_DIR, f"state_{int(time.time()*1000)}_{random.randint(0,9999):04d}.npz")
|
|
np.savez_compressed(state_path, **state_data)
|
|
return state_path
|
|
|
|
def unpack_state(state_path):
|
|
data = np.load(state_path)
|
|
shape_slat = SparseTensor(
|
|
feats=torch.from_numpy(data['shape_slat_feats']).cuda(),
|
|
coords=torch.from_numpy(data['coords']).cuda(),
|
|
)
|
|
tex_slat = shape_slat.replace(torch.from_numpy(data['tex_slat_feats']).cuda())
|
|
return shape_slat, tex_slat, int(data['res'])
|
|
|
|
# ============================================================================
|
|
# Progress Tracking (file-based, cross-process safe for @spaces.GPU)
|
|
# ============================================================================
|
|
|
|
import asyncio
|
|
from fastapi.responses import JSONResponse
|
|
from fastapi import Request
|
|
|
|
PROGRESS_DIR = os.path.join(TMP_DIR, '_progress')
|
|
os.makedirs(PROGRESS_DIR, exist_ok=True)
|
|
|
|
_thread_local = threading.local()
|
|
|
|
def _progress_file(session_id: str) -> str:
|
|
"""Return path to a session's progress JSON file."""
|
|
return os.path.join(PROGRESS_DIR, f"{session_id}.json")
|
|
|
|
def _reset_progress(session_id: str):
|
|
_thread_local.active_session = session_id
|
|
_write_progress_file(session_id, {"stage": "Initializing...", "step": 0, "total": 0, "done": False})
|
|
|
|
def _update_progress(stage: str, step: int, total: int):
|
|
session_id = getattr(_thread_local, 'active_session', '')
|
|
if session_id:
|
|
_write_progress_file(session_id, {"stage": stage, "step": step, "total": total, "done": False})
|
|
|
|
def _finish_progress():
|
|
session_id = getattr(_thread_local, 'active_session', '')
|
|
if session_id:
|
|
_write_progress_file(session_id, {"done": True})
|
|
|
|
def _write_progress_file(session_id: str, data: dict):
|
|
"""Atomically write progress JSON to a file (cross-process safe)."""
|
|
path = _progress_file(session_id)
|
|
tmp_path = path + ".tmp"
|
|
try:
|
|
with open(tmp_path, 'w') as f:
|
|
json.dump(data, f)
|
|
os.replace(tmp_path, path) # atomic on POSIX
|
|
except Exception:
|
|
pass
|
|
|
|
# Monkey-patch tqdm to intercept progress
|
|
import tqdm as _tqdm_module
|
|
|
|
_original_tqdm = _tqdm_module.tqdm
|
|
|
|
class _TqdmProgressInterceptor(_original_tqdm):
|
|
"""Wraps tqdm to push progress updates to SSE."""
|
|
def __init__(self, *args, **kwargs):
|
|
self._stage_desc = kwargs.get('desc', 'Processing')
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def set_description(self, desc=None, refresh=True):
|
|
self._stage_desc = desc or 'Processing'
|
|
super().set_description(desc, refresh)
|
|
|
|
def update(self, n=1):
|
|
super().update(n)
|
|
_update_progress(self._stage_desc, self.n, self.total or 0)
|
|
|
|
# Patch tqdm globally
|
|
_tqdm_module.tqdm = _TqdmProgressInterceptor
|
|
# Also patch the direct import in the sampler module and render_utils
|
|
import pixal3d.pipelines.samplers.flow_euler as _fe_module
|
|
_fe_module.tqdm = _TqdmProgressInterceptor
|
|
import pixal3d.utils.render_utils as _ru_module
|
|
_ru_module.tqdm = _TqdmProgressInterceptor
|
|
import o_voxel.postprocess as _ovp_module
|
|
_ovp_module.tqdm = _TqdmProgressInterceptor
|
|
|
|
# ============================================================================
|
|
# API Implementation
|
|
# ============================================================================
|
|
|
|
app = Server()
|
|
|
|
@app.get("/")
|
|
async def homepage():
|
|
html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
|
|
with open(html_path, "r", encoding="utf-8") as f:
|
|
return HTMLResponse(content=f.read())
|
|
|
|
@app.get("/app_config")
|
|
async def get_config():
|
|
"""Return server configuration for frontend (e.g. LOW_VRAM mode)."""
|
|
return JSONResponse({"low_vram": LOW_VRAM})
|
|
|
|
@app.get("/progress")
|
|
async def progress_poll(request: Request):
|
|
"""Polling endpoint for real-time progress updates during generation."""
|
|
session_id = request.query_params.get("session_id", "")
|
|
path = _progress_file(session_id)
|
|
try:
|
|
with open(path, 'r') as f:
|
|
data = json.load(f)
|
|
return JSONResponse(data)
|
|
except (FileNotFoundError, json.JSONDecodeError):
|
|
return JSONResponse({"stage": "Waiting...", "step": 0, "total": 0, "done": False})
|
|
|
|
@app.api()
|
|
@spaces.GPU(duration=30)
|
|
def preprocess(image: FileData) -> FileData:
|
|
init_models()
|
|
img = Image.open(image["path"])
|
|
processed = pipeline.preprocess_image(img)
|
|
out_path = os.path.join(TMP_DIR, f"preprocessed_{int(time.time()*1000)}.png")
|
|
processed.save(out_path)
|
|
return FileData(path=out_path)
|
|
|
|
@app.api()
|
|
@spaces.GPU(duration=120)
|
|
def generate_3d(
|
|
image: FileData,
|
|
seed: int,
|
|
resolution: int,
|
|
ss_guidance_strength: float = 7.5,
|
|
ss_guidance_rescale: float = 0.7,
|
|
ss_sampling_steps: int = 12,
|
|
ss_rescale_t: float = 5.0,
|
|
shape_slat_guidance_strength: float = 7.5,
|
|
shape_slat_guidance_rescale: float = 0.5,
|
|
shape_slat_sampling_steps: int = 12,
|
|
shape_slat_rescale_t: float = 3.0,
|
|
tex_slat_guidance_strength: float = 1.0,
|
|
tex_slat_guidance_rescale: float = 0.0,
|
|
tex_slat_sampling_steps: int = 12,
|
|
tex_slat_rescale_t: float = 3.0,
|
|
manual_fov: float = -1.0,
|
|
fov_unit: str = "deg",
|
|
session_id: str = "",
|
|
) -> Dict:
|
|
init_models()
|
|
_reset_progress(session_id)
|
|
_update_progress("Preprocessing & Camera Estimation", 0, 1)
|
|
|
|
torch.manual_seed(seed)
|
|
hr_resolution = int(resolution)
|
|
|
|
img = Image.open(image["path"])
|
|
# Image is already preprocessed by /preprocess endpoint, use directly
|
|
image_preprocessed = img
|
|
temp_processed_path = os.path.join(TMP_DIR, f"temp_proc_{session_id[:8]}_{int(time.time()*1000)}.png")
|
|
image_preprocessed.save(temp_processed_path)
|
|
|
|
if manual_fov > 0:
|
|
# Convert to radians based on unit
|
|
if fov_unit == "rad":
|
|
camera_angle_x = float(manual_fov)
|
|
fov_deg = math.degrees(manual_fov)
|
|
else:
|
|
camera_angle_x = math.radians(manual_fov)
|
|
fov_deg = float(manual_fov)
|
|
grid_point = torch.tensor([-1.0, 0.0, 0.0])
|
|
distance = distance_from_fov(
|
|
camera_angle_x, grid_point,
|
|
torch.tensor([0 - WILD_EXTEND_PIXEL, WILD_IMAGE_RESOLUTION - 1 + WILD_EXTEND_PIXEL]),
|
|
WILD_MESH_SCALE, WILD_IMAGE_RESOLUTION
|
|
)["distance_from_x"]
|
|
camera_params = {'camera_angle_x': camera_angle_x, 'distance': distance, 'mesh_scale': WILD_MESH_SCALE}
|
|
print(f"[Camera] Using manual FOV: {fov_deg:.2f}° ({camera_angle_x:.4f} rad), distance: {distance:.4f}")
|
|
else:
|
|
camera_params = get_camera_params_wild_moge(
|
|
temp_processed_path, device="cuda",
|
|
mesh_scale=WILD_MESH_SCALE, extend_pixel=WILD_EXTEND_PIXEL,
|
|
image_resolution=WILD_IMAGE_RESOLUTION,
|
|
)
|
|
_update_progress("Preprocessing & Camera Estimation", 1, 1)
|
|
|
|
ss_sampler_override = {"steps": ss_sampling_steps, "guidance_strength": ss_guidance_strength,
|
|
"guidance_rescale": ss_guidance_rescale, "rescale_t": ss_rescale_t}
|
|
shape_sampler_override = {"steps": shape_slat_sampling_steps, "guidance_strength": shape_slat_guidance_strength,
|
|
"guidance_rescale": shape_slat_guidance_rescale, "rescale_t": shape_slat_rescale_t}
|
|
tex_sampler_override = {"steps": tex_slat_sampling_steps, "guidance_strength": tex_slat_guidance_strength,
|
|
"guidance_rescale": tex_slat_guidance_rescale, "rescale_t": tex_slat_rescale_t}
|
|
|
|
pipeline_type = f"{hr_resolution}_cascade"
|
|
mesh_list, (shape_slat, tex_slat, res) = pipeline.run(
|
|
image_preprocessed,
|
|
camera_params=camera_params,
|
|
seed=seed,
|
|
sparse_structure_sampler_params=ss_sampler_override,
|
|
shape_slat_sampler_params=shape_sampler_override,
|
|
tex_slat_sampler_params=tex_sampler_override,
|
|
preprocess_image=False,
|
|
return_latent=True,
|
|
pipeline_type=pipeline_type,
|
|
max_num_tokens=CASCADE_MAX_NUM_TOKENS,
|
|
)
|
|
|
|
mesh = mesh_list[0]
|
|
state_path = pack_state(shape_slat, tex_slat, res)
|
|
|
|
_update_progress("Rendering views", 0, 1)
|
|
mesh.simplify(16777216)
|
|
cam_dist = camera_params['distance']
|
|
near = max(0.01, cam_dist - 2.0)
|
|
far = cam_dist + 10.0
|
|
if LOW_VRAM:
|
|
for v in envmap.values():
|
|
v.image = v.image.cuda()
|
|
if hasattr(v, '_nvdiffrec_envlight'):
|
|
del v._nvdiffrec_envlight
|
|
renders = render_utils.render_proj_aligned_video(
|
|
mesh, camera_angle_x=camera_params['camera_angle_x'],
|
|
distance=cam_dist, resolution=1024,
|
|
num_frames=STEPS, envmap=envmap,
|
|
near=near, far=far,
|
|
)
|
|
if LOW_VRAM:
|
|
for v in envmap.values():
|
|
if hasattr(v, '_nvdiffrec_envlight'):
|
|
del v._nvdiffrec_envlight
|
|
v.image = v.image.cpu()
|
|
torch.cuda.empty_cache()
|
|
_update_progress("Rendering views", 1, 1)
|
|
|
|
# Save renders and return paths
|
|
render_files = {}
|
|
for mode_key, frames in renders.items():
|
|
mode_files = []
|
|
for i, frame in enumerate(frames):
|
|
p = os.path.abspath(os.path.join(TMP_DIR, f"render_{mode_key}_{i}_{int(time.time()*1000)}.jpg"))
|
|
Image.fromarray(frame).save(p, quality=85)
|
|
mode_files.append(FileData(path=p))
|
|
render_files[mode_key] = mode_files
|
|
|
|
_finish_progress()
|
|
return {
|
|
"render_paths": render_files,
|
|
"state_path": os.path.abspath(state_path),
|
|
"camera_angle_x": camera_params['camera_angle_x'],
|
|
"distance": camera_params['distance'],
|
|
}
|
|
|
|
@app.api()
|
|
@spaces.GPU(duration=240)
|
|
def extract_glb_api(state_path: str, decimation_target: int, texture_size: int, session_id: str = "") -> FileData:
|
|
init_models()
|
|
_reset_progress(session_id)
|
|
_update_progress("Decoding latent", 0, 1)
|
|
|
|
shape_slat, tex_slat, res = unpack_state(state_path)
|
|
mesh = pipeline.decode_latent(shape_slat, tex_slat, res)[0]
|
|
_update_progress("Decoding latent", 1, 1)
|
|
|
|
glb = o_voxel.postprocess.to_glb(
|
|
vertices=mesh.vertices, faces=mesh.faces, attr_volume=mesh.attrs,
|
|
coords=mesh.coords, attr_layout=pipeline.pbr_attr_layout,
|
|
grid_size=res, 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,
|
|
)
|
|
rot = np.array([
|
|
[-1, 0, 0, 0],
|
|
[ 0, 0, -1, 0],
|
|
[ 0, -1, 0, 0],
|
|
[ 0, 0, 0, 1],
|
|
], dtype=np.float64)
|
|
glb.apply_transform(rot)
|
|
|
|
out_glb = os.path.join(TMP_DIR, f"result_{int(time.time()*1000)}.glb")
|
|
glb.export(out_glb, extension_webp=True)
|
|
_finish_progress()
|
|
return FileData(path=out_glb)
|
|
|
|
# Mount assets and tmp for direct access
|
|
app.mount("/assets", StaticFiles(directory="assets"), name="assets")
|
|
app.mount("/tmp", StaticFiles(directory=TMP_DIR), name="tmp")
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
parser = argparse.ArgumentParser(description="Pixal3D Demo Server")
|
|
parser.add_argument("--low_vram", action="store_true",
|
|
help="Enable low-VRAM mode: models lazy-load to GPU per stage.")
|
|
args, remaining = parser.parse_known_args()
|
|
if args.low_vram:
|
|
LOW_VRAM = True
|
|
|
|
# Re-install utils3d as in original app.py
|
|
subprocess.run([
|
|
sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps",
|
|
"https://github.com/LDYang694/Storages/releases/download/20260430/utils3d-0.0.2-py3-none-any.whl"
|
|
], check=True)
|
|
|
|
# Pre-initialize models before launching the server
|
|
init_models()
|
|
|
|
app.launch(show_error=True, share=True) |