Spaces:
Running on Zero
Running on Zero
| """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} | |