ADRIEN98's picture
Create app.py
9602d4e verified
import torch
from diffusers import DiffusionPipeline
import gradio as gr
import os
# 1. Load the Model
# Use 'cuda' if you have an NVIDIA GPU and CUDA installed, otherwise 'cpu'
# Hugging Face Spaces will handle the GPU for you if you selected it
device = "cuda" if torch.cuda.is_available() else "cpu"
try:
pipe = DiffusionPipeline.from_pretrained("cerspense/zeroscope_v2_576w", torch_dtype=torch.float16 if device == "cuda" else torch.float32)
pipe.to(device)
print(f"Model loaded successfully on {device}!")
except Exception as e:
print(f"Error loading model: {e}")
print("Ensure all dependencies are in requirements.txt and GPU is available if using 'cuda'.")
pipe = None # Set pipe to None if loading fails
# 2. Define the Generation Function
def generate_video(prompt: str, num_frames: int = 24, guidance_scale: float = 9.0):
if pipe is None:
return "Error: AI model not loaded. Please check server logs in the 'Logs' tab."
if not prompt:
return "Please enter a text prompt."
print(f"Generating video for prompt: '{prompt}'")
try:
# Generate video frames
video_frames = pipe(
prompt=prompt,
num_frames=num_frames,
guidance_scale=guidance_scale
).frames
# The output file path for the video
output_filename = "generated_video.mp4"
# pipe.to_video creates the video file from frames
video_path = pipe.to_video(video_frames, output_filename)
return video_path
except Exception as e:
print(f"Video generation failed: {e}")
return f"Video generation failed: {e}"
# 3. Create Gradio Interface
iface = gr.Interface(
fn=generate_video,
inputs=[
gr.Textbox(label="Text Prompt", placeholder="e.g., 'A robot dancing in the rain'"),
gr.Slider(minimum=16, maximum=64, step=4, value=24, label="Number of Frames (Video Length)"),
gr.Slider(minimum=1.0, maximum=15.0, step=0.5, value=9.0, label="Guidance Scale (Creativity vs. Prompt Adherence)")
],
outputs=gr.Video(label="Generated Video"),
title="Simple Text-to-Video Generator (Zeroscope v2)",
description="Enter a text prompt to generate a short video using Zeroscope v2."
)
# 4. Launch the App
if __name__ == "__main__":
iface.launch() # Hugging Face Spaces handles sharing automatically