Spaces:
Sleeping
Sleeping
| import spaces | |
| import os | |
| import uuid | |
| import tempfile | |
| import torch | |
| import gradio as gr | |
| from diffusers import ( | |
| AutoencoderKLWan, | |
| WanPipeline | |
| ) | |
| from diffusers.schedulers.scheduling_unipc_multistep import ( | |
| UniPCMultistepScheduler | |
| ) | |
| from diffusers.utils import export_to_video | |
| # ============================================================ | |
| # CONFIGURATION | |
| # ============================================================ | |
| MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" | |
| OUTPUT_DIR = "/tmp/learnmate_videos" | |
| os.makedirs(OUTPUT_DIR, exist_ok=True) | |
| # ============================================================ | |
| # STARTUP INFORMATION | |
| # ============================================================ | |
| print("=" * 60) | |
| print("LearnMate Video Generator") | |
| print("=" * 60) | |
| print("Model:", MODEL_ID) | |
| print("PyTorch:", torch.__version__) | |
| print("CUDA available:", | |
| torch.cuda.is_available()) | |
| # ============================================================ | |
| # LOAD WAN MODEL | |
| # ============================================================ | |
| # | |
| # IMPORTANT: | |
| # ZeroGPU requires the model to be placed on CUDA | |
| # at module level. | |
| # | |
| # Hugging Face ZeroGPU emulates CUDA during startup, | |
| # and provides the real GPU when @spaces.GPU is called. | |
| # ============================================================ | |
| print("Loading Wan VAE...") | |
| vae = AutoencoderKLWan.from_pretrained( | |
| MODEL_ID, | |
| subfolder="vae", | |
| torch_dtype=torch.float32 | |
| ) | |
| print("Loading Wan pipeline...") | |
| pipe = WanPipeline.from_pretrained( | |
| MODEL_ID, | |
| vae=vae, | |
| torch_dtype=torch.bfloat16 | |
| ) | |
| # 480p recommended flow shift | |
| pipe.scheduler = UniPCMultistepScheduler.from_config( | |
| pipe.scheduler.config, | |
| flow_shift=3.0 | |
| ) | |
| # Move model to CUDA | |
| pipe.to("cuda") | |
| print("Wan model loaded successfully!") | |
| print("=" * 60) | |
| # ============================================================ | |
| # VIDEO GENERATION | |
| # ============================================================ | |
| def generate_video(prompt): | |
| if not prompt or not prompt.strip(): | |
| raise gr.Error( | |
| "Please enter a video description." | |
| ) | |
| print("=" * 60) | |
| print("Generating video...") | |
| print("Prompt:", prompt) | |
| # -------------------------------------------------------- | |
| # Negative prompt | |
| # -------------------------------------------------------- | |
| negative_prompt = """ | |
| blurry, | |
| low quality, | |
| distorted, | |
| static image, | |
| still picture, | |
| subtitles, | |
| text, | |
| watermark, | |
| logo, | |
| deformed, | |
| disfigured, | |
| ugly, | |
| extra limbs, | |
| extra fingers, | |
| poorly drawn hands, | |
| poorly drawn faces, | |
| messy background, | |
| duplicate objects | |
| """ | |
| # -------------------------------------------------------- | |
| # Generate video | |
| # -------------------------------------------------------- | |
| try: | |
| output = pipe( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| # 480p | |
| height=480, | |
| width=832, | |
| # Short video | |
| num_frames=33, | |
| # Lower number = faster | |
| num_inference_steps=20, | |
| guidance_scale=5.0 | |
| ) | |
| frames = output.frames[0] | |
| # ---------------------------------------------------- | |
| # Save MP4 | |
| # ---------------------------------------------------- | |
| video_id = str(uuid.uuid4()) | |
| video_path = os.path.join( | |
| OUTPUT_DIR, | |
| f"{video_id}.mp4" | |
| ) | |
| export_to_video( | |
| frames, | |
| video_path, | |
| fps=16 | |
| ) | |
| print("Video generated:") | |
| print(video_path) | |
| print("=" * 60) | |
| return video_path | |
| except Exception as e: | |
| print("VIDEO GENERATION ERROR:") | |
| print(str(e)) | |
| raise gr.Error( | |
| f"Video generation failed: {str(e)}" | |
| ) | |
| # ============================================================ | |
| # GRADIO INTERFACE | |
| # ============================================================ | |
| description = """ | |
| Generate short educational animations from a text description | |
| using Wan2.1 T2V 1.3B. | |
| Designed for the LearnMate personalized learning platform. | |
| """ | |
| examples = [ | |
| [ | |
| "Educational animation showing photosynthesis. " | |
| "A green plant absorbs sunlight through its leaves, " | |
| "takes water from the soil through its roots, " | |
| "and releases oxygen." | |
| ], | |
| [ | |
| "Simple educational animation explaining the water cycle. " | |
| "Water evaporates from the ocean, forms clouds, " | |
| "then falls as rain and returns to rivers." | |
| ], | |
| [ | |
| "Educational animation showing the solar system. " | |
| "The Earth slowly revolves around the Sun " | |
| "while the planets move in their orbits." | |
| ] | |
| ] | |
| demo = gr.Interface( | |
| fn=generate_video, | |
| inputs=gr.Textbox( | |
| label="Learning Video Prompt", | |
| placeholder="Describe the educational animation...", | |
| lines=5 | |
| ), | |
| outputs=gr.Video( | |
| label="Generated Learning Video", | |
| format="mp4", | |
| autoplay=False | |
| ), | |
| title="LearnMate 🎓 → 🎬", | |
| description=description, | |
| examples=examples, | |
| api_name="generate_video" | |
| ) | |
| # ============================================================ | |
| # LAUNCH | |
| # ============================================================ | |
| demo.launch() |