raw
history blame
1.17 kB
import torch,os
from diffusers import StableVideoDiffusionPipeline
from diffusers.utils import load_image, export_to_video
from PIL import Image
from huggingface_hub import login
login(token=os.getenv("TOKEN"))
# Check if CUDA (GPU) is available, otherwise use CPU
device = "cuda" if torch.cuda.is_available() else "cpu"
# Function to generate the video
def Video(image):
pipeline = StableVideoDiffusionPipeline.from_pretrained(
"stabilityai/stable-video-diffusion-img2vid-xt-1-1", torch_dtype=torch.float32
).to(device)
# Enable model offloading if using the CPU
if device == "cpu":
pipeline.enable_model_cpu_offload()
image = Image.fromarray(image)
image = image.resize((1024, 576))
# Set random seed for reproducibility
generator = torch.manual_seed(42)
# Ensure the image is moved to the appropriate device (GPU or CPU)
# image = image.to(device)
# Generate the video frames
frames = pipeline(image, decode_chunk_size=8, generator=generator).frames[0]
# Export the frames to a video file
export_to_video(frames, "generated.mp4", fps=7)
return "generated.mp4"