File size: 2,194 Bytes
76f7b71 |
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 |
import os
import gradio as gr
from diffusers import DiffusionPipeline
import torch
# === Configure cache directory ===
cache_dir = os.path.expanduser("~/Downloads/Openking")
os.makedirs(cache_dir, exist_ok=True)
# Set Hugging Face cache environment variables
os.environ["HF_HOME"] = cache_dir
os.environ["HF_HUB_CACHE"] = cache_dir
os.environ["HF_DATASETS_CACHE"] = cache_dir
# === Load Hugging Face token from secrets (required for private models) ===
# In Hugging Face Spaces, store your token as a secret named "HF_TOKEN"
hf_token = os.getenv("HF_TOKEN")
if not hf_token:
raise ValueError("Please set your Hugging Face token as a secret named 'HF_TOKEN' in your Space settings.")
# === Load the model ===
model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
try:
pipe = DiffusionPipeline.from_pretrained(
model_id,
use_auth_token=hf_token,
cache_dir=cache_dir,
torch_dtype=torch.float16,
variant="fp16"
)
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
except Exception as e:
raise RuntimeError(f"Failed to load model: {e}")
# === Gradio interface ===
def generate_video(prompt: str, num_inference_steps: int = 50):
try:
# Note: Adjust this call based on the actual model's inference API.
# Since this is a text-to-video model, the exact method may vary.
# This is a placeholder—check the model card for correct usage.
video_frames = pipe(prompt, num_inference_steps=num_inference_steps).frames
# For now, return a placeholder message
return f"Generated video for: '{prompt}' with {num_inference_steps} steps. (Output handling depends on model output format.)"
except Exception as e:
return f"Error: {str(e)}"
with gr.Blocks() as demo:
gr.Markdown("# 🎥 Wan2.1 Text-to-Video Generator")
prompt = gr.Textbox(label="Prompt", placeholder="A cat flying through space...")
steps = gr.Slider(10, 100, value=50, label="Inference Steps")
output = gr.Textbox(label="Result")
btn = gr.Button("Generate Video")
btn.click(generate_video, inputs=[prompt, steps], outputs=output)
# Launch app
if __name__ == "__main__":
demo.launch() |