Spaces:
Runtime error
Runtime error
| """ZeroGPU workers inherit startup models; only CPU inputs/results cross their process queues.""" | |
| import json | |
| import math | |
| import os | |
| from pathlib import Path | |
| import sys | |
| import spaces | |
| ROOT = Path(__file__).resolve().parents[1] | |
| RUNTIME = ROOT / ".runtime" | |
| LOCK = json.loads((ROOT / "models.lock.json").read_text()) | |
| _MODELS = {} | |
| def reference_gpu(image, prompt, seed): | |
| import torch | |
| kwargs = {"prompt": prompt + ". Single complete asset, full silhouette visible, plain background, no text.", | |
| "height": 1024, "width": 1024, "num_inference_steps": 4, "guidance_scale": 1.0, | |
| "generator": torch.Generator("cuda").manual_seed(seed)} | |
| if image is not None: | |
| kwargs["image"] = image.convert("RGB") | |
| return _MODELS["reference"](**kwargs).images[0] | |
| def geometry_gpu(image, seed, output): | |
| import o_voxel | |
| import torch | |
| pipeline = _MODELS["geometry"] | |
| with torch.inference_mode(): | |
| meshes, latents = pipeline.run( | |
| image, seed=seed, pipeline_type="1024_cascade", return_latent=True, | |
| sparse_structure_sampler_params={"steps": 12}, | |
| shape_slat_sampler_params={"steps": 12}, tex_slat_sampler_params={"steps": 12}) | |
| mesh = meshes[0] | |
| mesh.simplify(16777216) | |
| 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=latents[2], | |
| aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], decimation_target=50000, | |
| texture_size=1024, remesh=True, remesh_band=1, remesh_project=0, use_tqdm=False) | |
| glb.export(output) | |
| return output | |
| def segment_gpu(mesh, count, seed): | |
| from .partfield import segment | |
| return segment(_MODELS["segment"], mesh, count, seed) | |
| def upscale_duration(parts): | |
| tiles = sum(math.ceil(part.visual.material.baseColorTexture.width / 256) * | |
| math.ceil(part.visual.material.baseColorTexture.height / 256) for _, part in parts) | |
| return min(180, max(30, math.ceil(15 + tiles * 1.5))) | |
| def upscale_gpu(parts): | |
| from .upscale import upscale_parts | |
| events = [] | |
| stats = upscale_parts(parts, _MODELS["sr"], events.append) | |
| textures = {key: part.visual.material.baseColorTexture for key, part in parts} | |
| return textures, stats, events | |
| class NativeModels: | |
| def __init__(self, role=None): | |
| from scripts.prepare_runtime import worker_role | |
| self.role = worker_role(role) | |
| marker = "prepared.json" if self.role == "all" else f"prepared-{self.role}.json" | |
| if not (RUNTIME / marker).exists(): | |
| raise RuntimeError(f"Prepare the {self.role} worker runtime before loading GPU models.") | |
| if self.role in ("all", "reference"): | |
| import torch | |
| from diffusers import Flux2KleinPipeline | |
| _MODELS["reference"] = Flux2KleinPipeline.from_pretrained( | |
| "black-forest-labs/FLUX.2-klein-4B", revision=LOCK["black-forest-labs/FLUX.2-klein-4B"], | |
| torch_dtype=torch.bfloat16).to("cuda") | |
| if self.role in ("all", "geometry"): | |
| sys.path.insert(0, str(RUNTIME / "trellis")) | |
| os.environ.setdefault("ATTN_BACKEND", "flash_attn") | |
| from trellis2.pipelines import Trellis2ImageTo3DPipeline | |
| pipeline = Trellis2ImageTo3DPipeline.from_pretrained(str(RUNTIME / "trellis-model")) | |
| pipeline.low_vram = True | |
| pipeline.cuda() | |
| _MODELS["geometry"] = pipeline | |
| if self.role in ("all", "segment"): | |
| sys.path.insert(0, str(RUNTIME / "PartField")) | |
| from huggingface_hub import hf_hub_download | |
| from .partfield import load_model | |
| checkpoint = hf_hub_download("mikaelaangel/partfield-ckpt", "model_objaverse.ckpt", | |
| revision=LOCK["mikaelaangel/partfield-ckpt"]) | |
| _MODELS["segment"] = load_model(RUNTIME / "PartField", checkpoint) | |
| if self.role in ("all", "texture"): | |
| from spandrel import ModelLoader | |
| _MODELS["sr"] = ModelLoader().load_from_file(str(RUNTIME / "RealESRGAN_x4plus.pth")).eval().cuda() | |
| if self.role in ("all", "motion"): | |
| from .motion import load_motion_model | |
| load_motion_model(RUNTIME, LOCK) | |
| # Static functions avoid serializing model instances as a `self` argument to ZeroGPU workers. | |
| reference = staticmethod(reference_gpu) | |
| geometry = staticmethod(geometry_gpu) | |
| segment = staticmethod(segment_gpu) | |
| from .motion import generate_motion | |
| motion = staticmethod(generate_motion) | |
| def upscale(self, parts, log): | |
| textures, stats, events = upscale_gpu(parts) | |
| # Worker mutations do not propagate to the parent process. Apply returned CPU images explicitly. | |
| for key, part in parts: | |
| part.visual.material.baseColorTexture = textures[key] | |
| for event in events: | |
| log(event) | |
| return stats | |
| def compare_geometry_gpu(image, seed, output): | |
| import torch | |
| from scripts.compare_geometry_exports import capture, compare | |
| pipeline = _MODELS["geometry"] | |
| with torch.inference_mode(): | |
| meshes, latents = pipeline.run( | |
| image, seed=seed, pipeline_type="1024_cascade", return_latent=True, | |
| sparse_structure_sampler_params={"steps": 12}, | |
| shape_slat_sampler_params={"steps": 12}, tex_slat_sampler_params={"steps": 12}) | |
| capture(meshes[0], latents[2], pipeline.pbr_attr_layout, output) | |
| compare(output) | |
| return output | |
| def replay_geometry_gpu(raw, layout, output, filtered=None): | |
| from scripts.replay_geometry_export import replay | |
| return replay(raw, layout, output, filtered) | |
| def geometry_1536_gpu(image, seed, output): | |
| import time | |
| import torch | |
| from scripts.compare_geometry_exports import capture | |
| import o_voxel | |
| pipeline = _MODELS["geometry"] | |
| output = Path(output) | |
| conditioned = pipeline.preprocess_image(image) | |
| conditioned.save(output / "conditioning.png") | |
| started = time.monotonic() | |
| torch.cuda.reset_peak_memory_stats() | |
| with torch.inference_mode(): | |
| meshes, latents = pipeline.run( | |
| conditioned, seed=seed, pipeline_type="1536_cascade", preprocess_image=False, | |
| return_latent=True, max_num_tokens=49152, | |
| sparse_structure_sampler_params={"steps": 12}, | |
| shape_slat_sampler_params={"steps": 12}, tex_slat_sampler_params={"steps": 12}) | |
| mesh = meshes[0] | |
| capture(mesh, latents[2], pipeline.pbr_attr_layout, output) | |
| stats = {"requested_resolution":1536, "effective_resolution":int(latents[2]), | |
| "tokens":int(latents[0].coords.shape[0]), "seed":int(seed), | |
| "generation_seconds":time.monotonic()-started, | |
| "peak_cuda_allocated":torch.cuda.max_memory_allocated(), | |
| "peak_cuda_reserved":torch.cuda.max_memory_reserved()} | |
| (output / "generation.json").write_text(json.dumps(stats,indent=2)) | |
| 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=latents[2], | |
| aabb=[[-.5,-.5,-.5],[.5,.5,.5]],decimation_target=50000,texture_size=1024, | |
| remesh=True,remesh_band=1,remesh_project=0,use_tqdm=False) | |
| glb.export(output / "current-remesh.glb") | |
| return str(output) | |