Spaces:
Running on Zero
Running on Zero
File size: 5,127 Bytes
5c331a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """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 = {}
@spaces.GPU(duration=60)
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]
@spaces.GPU(duration=180)
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
@spaces.GPU(duration=90)
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)))
@spaces.GPU(duration=upscale_duration)
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
|