Spaces:
Running on Zero
Running on Zero
File size: 3,498 Bytes
e0177dc | 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 126 127 128 129 130 131 132 133 | """Hugging Face Space entry point for InstructAV2AV."""
from __future__ import annotations
import logging
import os
import tempfile
from pathlib import Path
from types import SimpleNamespace
from typing import Any
import gradio as gr
try:
import spaces
except ImportError:
class _SpacesFallback:
@staticmethod
def GPU(*decorator_args: Any, **_decorator_kwargs: Any):
if decorator_args and callable(decorator_args[0]):
return decorator_args[0]
return lambda function: function
spaces = _SpacesFallback()
from scripts.demo import CSS, DEFAULT_CONFIG, DemoRuntime, build_demo
from space_model_store import HubModelStore
logging.basicConfig(
level=logging.INFO,
format="[%(asctime)s] %(levelname)s: %(message)s",
)
def default_hf_home() -> Path:
persistent_storage = Path("/data")
if persistent_storage.is_dir() and os.access(persistent_storage, os.W_OK):
return persistent_storage / ".huggingface"
return Path.home() / ".cache" / "huggingface"
HF_HOME = Path(
os.getenv("HF_HOME", default_hf_home())
).expanduser()
MODEL_HOME = Path(
os.getenv("INSTRUCTAV2AV_MODEL_HOME", HF_HOME / "instructav2av")
).expanduser()
OUTPUT_DIR = Path(
os.getenv(
"INSTRUCTAV2AV_OUTPUT_DIR",
Path(tempfile.gettempdir()) / "instructav2av_outputs",
)
).expanduser()
GPU_DURATION = int(os.getenv("INSTRUCTAV2AV_ZEROGPU_DURATION", "300"))
model_store = HubModelStore(MODEL_HOME, hub_cache_dir=HF_HOME / "hub")
if os.getenv("INSTRUCTAV2AV_EAGER_DOWNLOAD", "1").lower() not in {
"0",
"false",
"no",
}:
try:
model_store.preload_default()
logging.info("Shared weights and the General checkpoint are cached.")
except Exception:
logging.exception(
"Startup model download failed; the first edit request will retry it."
)
runtime_args = SimpleNamespace(
config_file=str(DEFAULT_CONFIG),
model_dir=str(model_store.ckpt_dir / "InstructAV2AV"),
ckpt_dir=str(model_store.ckpt_dir),
output_dir=str(OUTPUT_DIR),
device=int(os.getenv("INSTRUCTAV2AV_CUDA_DEVICE", "0")),
no_cpu_offload=os.getenv("INSTRUCTAV2AV_CPU_OFFLOAD", "0").lower()
in {"0", "false", "no"},
)
runtime = DemoRuntime(
runtime_args,
checkpoint_resolver=model_store.resolve_checkpoint,
)
@spaces.GPU(duration=GPU_DURATION)
def zero_gpu_task(
operation: str,
video_value: Any,
instruction: str,
model_key: str,
seed: float,
sample_steps: float,
video_guidance_scale: float,
audio_guidance_scale: float,
progress: gr.Progress = gr.Progress(),
) -> str:
if operation == "warmup":
return runtime.warmup("general", progress)
if operation != "generate":
raise gr.Error(f"Unsupported ZeroGPU operation: {operation}")
return runtime.generate(
video_value,
instruction,
model_key,
seed,
sample_steps,
video_guidance_scale,
audio_guidance_scale,
progress,
)
demo = build_demo(
runtime,
zero_gpu_fn=zero_gpu_task,
)
demo.queue(max_size=8, default_concurrency_limit=1)
if __name__ == "__main__":
demo.launch(
server_name="0.0.0.0",
server_port=int(os.getenv("PORT", "7860")),
allowed_paths=[str(runtime.output_dir)],
show_error=True,
theme=gr.themes.Default(),
css=CSS,
max_file_size=os.getenv("GRADIO_MAX_FILE_SIZE", "500mb"),
)
|