Spaces:
Running on Zero
Running on Zero
File size: 5,671 Bytes
3718945 56cd1d4 4db2e43 56cd1d4 3718945 871867d 3718945 771c5cf 3718945 ee16d46 3718945 4db2e43 3718945 de64f6d 3718945 de64f6d 3718945 de64f6d 3718945 de64f6d 3718945 de64f6d 3718945 de64f6d 3718945 ee16d46 3718945 ee16d46 771c5cf 871867d 771c5cf 3718945 871867d 4db2e43 3718945 56cd1d4 | 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 121 122 123 124 125 | """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
@spaces.GPU(duration=40)
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)
@spaces.GPU(duration=40)
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)
@spaces.GPU(duration=scene_gpu_duration)
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)
|