"""Cloud-only stages: CPU validation, ZeroGPU sampling, then CPU exports.""" from pathlib import Path import json import os import random import tempfile import threading import time import numpy as np from backend import Engine from bootstrap import UPSTREAM_COMMIT from bvh_export import export_bvh from motion_io import FPS, normalize, read_motion, take_window from render import render_mp4 class CloudEngine(Engine): def __init__(self, source_root=None, output_root=None, *, catalog): # Do not instantiate the local catalog or resolve Windows dataset paths. cloud_root = Path(__file__).resolve().parent.parent self.root = Path(source_root or os.environ.get( "STYLE_SALAD_ROOT", str(cloud_root / "upstream"))).resolve() self.outputs = Path(output_root or os.environ.get( "SPACE_OUTPUT_DIR", str(cloud_root / "outputs"))).resolve() self.lock = threading.Lock() self.model = None self.catalog = catalog def eagerly_load(self, report=print): """Register frozen CUDA weights before Gradio starts on ZeroGPU. Inherited loading retains weights_only=True and the exact missing-key whitelist for the published trainable-delta checkpoint. """ with self.lock: self._load(report) import torch if self.device.type == "cuda": # ZeroGPU owns the real device mapping for each worker. self.device = torch.device("cuda") self.model.device = self.device self.model.opt.device = self.device self.model.vae_opt.device = self.device def prepare(self, style_path, caption, frames=120, seed=42, steps=34, text_weight=7.5, style_weight=1.5, style_guidance=0.75, content_path=None, report=lambda message: None, reference_info=None): """CPU-only validation; return plain values and small CPU numpy arrays.""" started = time.perf_counter() caption = str(caption or "").strip() if not caption or len(caption) > 300: raise ValueError("Describe the action in English, using at most 300 characters.") try: frames, seed, steps = int(frames), int(seed), int(steps) except (TypeError, ValueError, OverflowError) as error: raise ValueError("Frames, seed and diffusion steps must be integers.") from error if not 40 <= frames <= 196 or frames % 4: raise ValueError("Use 40 to 196 frames, in multiples of 4.") if not 0 <= seed <= 2**31 - 1 or not 5 <= steps <= 100: raise ValueError("Invalid seed or diffusion step count.") for value, upper in [(text_weight, 15), (style_weight, 5), (style_guidance, 2)]: if not np.isfinite(value) or not 0 <= value <= upper: raise ValueError("Guidance value outside the allowed range.") if self.model is None: raise RuntimeError("The cloud model must be loaded during startup.") report("Validating motion features...") style_raw = read_motion(style_path) content_raw = read_motion(content_path) if content_path else None drop_first = bool(self.config["dataset_style"].get("drop_first_frame", False)) style = take_window(style_raw, drop_first=drop_first) content = take_window(content_raw, max_frames=frames) if content_raw is not None else None length = content.frames if content is not None else frames return { "reference_normalized": normalize(style.features, self.mean, self.std), "reference_features": style.features, "content_normalized": None if content is None else normalize(content.features, self.mean, self.std), "content_features": None if content is None else content.features, "caption": caption, "frames": length, "requested_frames": frames, "seed": seed, "steps": steps, "text_weight": float(text_weight), "style_weight": float(style_weight), "style_guidance": float(style_guidance), "style_reference": reference_info or {"selection": "upload", "filename": Path(style_path).name}, "style_first_frame": style.first_frame, "style_frames": style.frames, "style_original_frames": style.original_frames, "style_drop_first_frame": drop_first, "content_first_frame": None if content is None else content.first_frame, "content_frames": None if content is None else content.frames, "content_original_frames": None if content is None else content.original_frames, "preparation_seconds": time.perf_counter() - started, } def infer(self, prepared, report=lambda message: None): """Call only from the registered @spaces.GPU handler; no rendering/I/O.""" if not isinstance(prepared, dict) or "reference_normalized" not in prepared: raise ValueError("Prepare a valid style reference before generating motion.") import torch if self.model is None: raise RuntimeError("The cloud model was not loaded during startup.") if self.device.type != "cuda" or not torch.cuda.is_available(): raise RuntimeError("A real ZeroGPU worker is required for motion inference.") with self.lock: seed, length = prepared["seed"], prepared["frames"] random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) ref = torch.from_numpy(prepared["reference_normalized"]).unsqueeze(0).to(self.device) ref_lengths = torch.tensor([prepared["style_frames"]], device=self.device, dtype=torch.long) lengths = torch.tensor([length], device=self.device, dtype=torch.long) content = prepared["content_normalized"] init_motion = None if content is None else torch.from_numpy(content).unsqueeze(0).to(self.device) self.model.config.update(text_weight=prepared["text_weight"], style_weight=prepared["style_weight"], style_guidance=prepared["style_guidance"]) guidance = {"num_inference_steps": prepared["steps"]} if content is not None: guidance["inversion"] = {"conditioning": "text"} report("Generating motion on ZeroGPU..." if content is None else "Inverting and stylizing the content motion on ZeroGPU...") torch.cuda.synchronize() started = time.perf_counter() try: # Frozen weights still support the derivative of style loss wrt z. with torch.enable_grad(): generated, _ = self.model.generate( ref, [prepared["caption"]], lengths, ref_lengths, guidance=guidance, init_motion=init_motion, init_lengths=lengths if content is not None else None, ) normalized = generated.detach()[0, :length].cpu().numpy() torch.cuda.synchronize() features = np.asarray(normalized * self.std + self.mean, dtype=np.float32) if features.shape != (length, 263) or not np.isfinite(features).all(): raise RuntimeError("Generation produced invalid features. Try a different reference or settings.") result = dict(prepared) result.update(features=features, inference_seconds=time.perf_counter() - started, gpu_model=torch.cuda.get_device_name(), torch_version=str(torch.__version__), cuda_version=str(torch.version.cuda), execution="huggingface-zerogpu") # Never return CUDA tensors across the ZeroGPU worker boundary. return result finally: self.model.denoiser.remove_clip_cache() self.model.denoiser.clear_hyper_lora_cache() torch.cuda.empty_cache() def finish(self, generated, report=lambda message: None): """CPU-only joint reconstruction, MP4, BVH, and download files.""" if not isinstance(generated, dict) or "features" not in generated: raise ValueError("Generate a motion before requesting its output files.") started = time.perf_counter() import torch features = np.asarray(generated["features"], dtype=np.float32) length = generated["frames"] if features.shape != (length, 263) or not np.isfinite(features).all(): raise ValueError("The generated motion features are invalid.") with torch.no_grad(): joints = self.recover(torch.from_numpy(features).unsqueeze(0), 22)[0].numpy() ref_joints = self.recover(torch.from_numpy(generated["reference_features"]).unsqueeze(0), 22)[0].numpy() content = generated["content_features"] content_joints = None if content is None else self.recover(torch.from_numpy(content).unsqueeze(0), 22)[0].numpy() if (not np.isfinite(joints).all() or not np.isfinite(ref_joints).all() or (content_joints is not None and not np.isfinite(content_joints).all())): raise RuntimeError("Motion reconstruction produced invalid joint positions.") self.outputs.mkdir(parents=True, exist_ok=True) folder = Path(tempfile.mkdtemp(prefix="motion-", dir=self.outputs)) report("Rendering videos on CPU...") video = render_mp4(joints, folder / "generated.mp4", FPS) reference = render_mp4(ref_joints, folder / "style_reference.mp4", FPS) content_video = None if content_joints is None else render_mp4(content_joints, folder / "content_reference.mp4", FPS) report("Exporting BVH on CPU...") bvh = export_bvh(joints, folder / "generated.bvh", FPS) np.save(folder / "generated_features.npy", features, allow_pickle=False) np.savez_compressed(folder / "generated_motion.npz", features=features, joints=joints, fps=np.array(FPS), caption=np.array(generated["caption"])) parameter_keys = ( "caption", "seed", "frames", "requested_frames", "steps", "text_weight", "style_weight", "style_guidance", "style_reference", "style_first_frame", "style_frames", "style_original_frames", "style_drop_first_frame", "content_first_frame", "content_frames", "content_original_frames", "preparation_seconds", "inference_seconds", "execution", "gpu_model", "torch_version", "cuda_version", ) info = {key: generated[key] for key in parameter_keys} render_seconds = time.perf_counter() - started info.update(upstream_commit=UPSTREAM_COMMIT, checkpoint="epoch_0100.ckpt", fps=FPS, mode="text_and_style" if content is None else "ddim_inversion_experimental", device="cuda", render_seconds=render_seconds, processing_seconds=info["preparation_seconds"] + info["inference_seconds"] + render_seconds, render="22-joint skeleton; camera follows root; world-fixed 0.5 m checkerboard and trail", mp4_source="generated features -> recover_from_ric -> joint positions (not BVH)", bvh={key: value for key, value in bvh.items() if key != "path"}) (folder / "metadata.json").write_text(json.dumps(info, indent=2, ensure_ascii=False), encoding="utf-8") return {"video": str(video), "reference": str(reference), "content_video": None if content_video is None else str(content_video), "features": str(folder / "generated_features.npy"), "motion": str(folder / "generated_motion.npz"), "bvh": bvh["path"], "metadata": str(folder / "metadata.json"), "info": info}