Spaces:
Running on Zero
Running on Zero
| """Shared ZeroGPU backend for the Object and Scene Hugging Face Spaces.""" | |
| import os | |
| from pathlib import Path | |
| from types import SimpleNamespace | |
| # rembg imports Numba. Its default TBB pool cannot safely fork from a | |
| # Gradio worker thread when ZeroGPU starts the GPU process. | |
| os.environ.setdefault("NUMBA_THREADING_LAYER", "workqueue") | |
| os.environ.setdefault("NUMBA_NUM_THREADS", "2") | |
| for variable in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS"): | |
| os.environ[variable] = "1" | |
| # ZeroGPU must install its CUDA emulation before torch is imported. | |
| import spaces | |
| import torch | |
| import gradio as gr | |
| from huggingface_hub import snapshot_download | |
| from omegaconf import OmegaConf | |
| from object_demo import ObjectDemo, build_demo as build_object_demo | |
| from scene_demo import SceneDemo, build_demo as build_scene_demo | |
| from diffusionGS.models.denoiser.denoiser import DGSDenoiser | |
| from diffusionGS.models.diffusion import create_diffusion | |
| from diffusionGS.pipeline_scene import DiffusionGSScenePipeline | |
| from diffusionGS.pipline_obj import DiffusionGSPipeline | |
| from space_success_stats import SuccessStats, UsageDisplay | |
| ROOT = Path(__file__).resolve().parent | |
| OUTPUT_ROOT = Path(os.environ.get("DIFFUSIONGS_OUTPUT_DIR", ROOT / "outputs")).resolve() | |
| MODEL_REPO = "CaiYuanhao/DiffusionGS" | |
| MODEL_REVISION = "6ad55569352ab32bd7ca52ba1c5f5cf89065b2d6" | |
| stats = SuccessStats.from_env() | |
| stats.start("diffusiongs-object", "diffusiongs-scene") | |
| def use_torch_attention(model): | |
| """Use PyTorch SDPA across the different GPUs allocated by ZeroGPU.""" | |
| for module in model.modules(): | |
| if hasattr(module, "use_flashatt_v2"): | |
| module.use_flashatt_v2 = False | |
| def load_object_pipeline(model_dir): | |
| # Inference only needs the denoiser, diffusion schedule and camera template. | |
| # Avoid allocating training losses and downloading their VGG checkpoint. | |
| config = OmegaConf.load(model_dir / "obj_configs.yaml") | |
| model = DGSDenoiser(OmegaConf.to_container(config.system.shape_model)) | |
| checkpoint = torch.load(model_dir / "obj_ckpt_512.ckpt", map_location="cpu", weights_only=False) | |
| state = checkpoint.get("state_dict", checkpoint) | |
| model.load_state_dict({key.removeprefix("shape_model."): value for key, value in state.items() | |
| if key.startswith("shape_model.")}, strict=True) | |
| del checkpoint, state | |
| use_torch_attention(model) | |
| model = model.to(device="cuda:0", dtype=torch.float16).eval() | |
| system = SimpleNamespace( | |
| shape_model=model, device=torch.device("cuda:0"), dtype=torch.float16, | |
| diffusion_inference=create_diffusion(str(config.system.num_inference_steps), predict_xstart=True), | |
| cam_template=torch.load(model_dir / "camera_template.pt", map_location="cpu", weights_only=False), | |
| ) | |
| return DiffusionGSPipeline("cuda:0", config, system) | |
| torch.set_num_threads(1) | |
| model_dir = Path(snapshot_download( | |
| MODEL_REPO, revision=MODEL_REVISION, | |
| allow_patterns=["obj_ckpt_512.ckpt", "obj_configs.yaml", "camera_template.pt", | |
| "scene_ckpt_256.ckpt", "scene_ckpt_512.ckpt"], | |
| )) | |
| object_service = ObjectDemo(output_dir=OUTPUT_ROOT / "object") | |
| object_service.pipeline = load_object_pipeline(model_dir) | |
| scene_pipelines = {} | |
| for resolution in (256, 512): | |
| pipeline = DiffusionGSScenePipeline.from_pretrained(model_dir, resolution=resolution) | |
| use_torch_attention(pipeline.model) | |
| scene_pipelines[resolution] = pipeline | |
| scene_service = SceneDemo(output_dir=OUTPUT_ROOT / "scene") | |
| _generate_object = object_service.generate | |
| _extract_mesh = object_service.extract_mesh | |
| _generate_scene = scene_service.generate | |
| def generate_object(image, seed, foreground_ratio, force_remove_background): | |
| print("Starting object generation on ZeroGPU", flush=True) | |
| return _generate_object(image, seed, foreground_ratio, force_remove_background) | |
| def extract_mesh(result_directory): | |
| print("Starting mesh extraction on ZeroGPU", flush=True) | |
| return _extract_mesh(result_directory) | |
| def scene_gpu_duration(image, preset, translation_scale, rotation_degrees, fov_degrees, | |
| duration_seconds, resolution, steps, seed, fps): | |
| return (20 + (20 if int(resolution) == 512 else 0) | |
| + max(0, int(steps) - 30) // 2 | |
| + max(0, int(duration_seconds) * int(fps) // 2 - 120) // 10) | |
| def generate_scene(image, preset, translation_scale, rotation_degrees, fov_degrees, | |
| duration_seconds, resolution, steps, seed, fps): | |
| print("Starting scene generation on ZeroGPU", flush=True) | |
| scene_service.pipeline = scene_pipelines[int(resolution)] | |
| return _generate_scene(image, preset, translation_scale, rotation_degrees, fov_degrees, | |
| duration_seconds, resolution, steps, seed, fps) | |
| object_service.generate = stats.successful("diffusiongs-object")(generate_object) | |
| object_service.extract_mesh = extract_mesh | |
| scene_service.generate = stats.successful("diffusiongs-scene")(generate_scene) | |
| demo = build_object_demo(service=object_service, usage=UsageDisplay(stats, "diffusiongs-object")) | |
| with demo: | |
| gr.Navbar(main_page_name="Object Generation") | |
| scene_demo = build_scene_demo(service=scene_service, usage=UsageDisplay(stats, "diffusiongs-scene")) | |
| with demo.route("Scene Generation", "/scene"): | |
| gr.Navbar(main_page_name="Object Generation") | |
| scene_demo.render() | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("GRADIO_SERVER_PORT", "7860")), | |
| max_file_size="20mb", allowed_paths=[str(OUTPUT_ROOT)], ssr_mode=False) | |