File size: 3,466 Bytes
14da46e e337195 14da46e e337195 14da46e e337195 14da46e 60b0713 14da46e e337195 14da46e e337195 14da46e e337195 14da46e e337195 14da46e e337195 14da46e | 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 | import os
import time
import gradio as gr
import torch
from PIL import Image
from huggingface_hub import snapshot_download
from diffusers import StableVideoDiffusionPipeline
import imageio.v2 as imageio
REPO_ID = "stabilityai/stable-video-diffusion-img2vid-xt"
# Где хранить файлы модели (пытаемся /data, если нет — /home/user)
ROOT = "/data" if os.path.isdir("/data") else "/home/user"
MODEL_DIR = os.path.join(ROOT, "models", "svd-xt")
OUT_DIR = "/tmp/outputs"
os.makedirs(OUT_DIR, exist_ok=True)
pipe = None
def ensure_model():
os.makedirs(MODEL_DIR, exist_ok=True)
# Если уже скачано — не качаем заново
if any(os.path.exists(os.path.join(MODEL_DIR, f)) for f in ["model_index.json", "config.json"]):
print("✅ Model already present in:", MODEL_DIR)
return
print("⬇️ Downloading model to:", MODEL_DIR)
snapshot_download(
repo_id=REPO_ID,
local_dir=MODEL_DIR,
local_dir_use_symlinks=False,
resume_download=True,
max_workers=4,
token=os.environ.get("HF_TOKEN"), # если модель gated
)
print("✅ Download finished. Top files:", os.listdir(MODEL_DIR)[:30])
def get_pipe():
global pipe
if pipe is not None:
return pipe
ensure_model()
dtype = torch.float16 if torch.cuda.is_available() else torch.float32
pipe = StableVideoDiffusionPipeline.from_pretrained(
MODEL_DIR,
torch_dtype=dtype,
local_files_only=True, # <-- важно: после скачивания не лезем в интернет
)
if torch.cuda.is_available():
pipe = pipe.to("cuda")
try:
pipe.enable_xformers_memory_efficient_attention()
except Exception:
pass
return pipe
def run(image: Image.Image, motion: int, fps: int, frames: int, steps: int, seed: int):
if image is None:
raise gr.Error("Upload an image first.")
if seed < 0:
seed = int(time.time()) % 10_000_000
generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(int(seed))
pipe = get_pipe()
img = image.convert("RGB")
out = pipe(
image=img,
num_frames=int(frames),
fps=int(fps),
motion_bucket_id=int(motion),
num_inference_steps=int(steps),
generator=generator,
)
video_frames = out.frames[0] # list[PIL]
out_path = os.path.join(OUT_DIR, f"svd_{int(time.time())}.mp4")
imageio.mimsave(out_path, video_frames, fps=int(fps))
return out_path
with gr.Blocks(title="SVD img2vid XT") as demo:
gr.Markdown("## Stable Video Diffusion (img2vid-xt)\nModel downloads once at startup (resume enabled).")
with gr.Row():
inp = gr.Image(type="pil", label="Input image")
out = gr.Video(label="Output video")
with gr.Accordion("Settings", open=False):
motion = gr.Slider(1, 255, value=100, step=1, label="motion_bucket_id (lower = calmer)")
fps = gr.Slider(6, 30, value=12, step=1, label="fps")
frames = gr.Slider(8, 30, value=14, step=1, label="num_frames")
steps = gr.Slider(10, 50, value=25, step=1, label="steps")
seed = gr.Number(value=-1, precision=0, label="seed (-1 random)")
btn = gr.Button("Generate", variant="primary")
btn.click(run, [inp, motion, fps, frames, steps, seed], out)
demo.queue().launch() |