Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from PIL import Image | |
| from rembg import remove | |
| import spaces | |
| import torch | |
| from diffusers import AutoencoderKLWan, WanImageToVideoPipeline | |
| from diffusers.utils import export_to_video | |
| from transformers import CLIPVisionModel | |
| MODEL_ID = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" | |
| print("Loading Wan2.1 I2V pipeline... (Using cache from last time)") | |
| image_encoder = CLIPVisionModel.from_pretrained(MODEL_ID, subfolder="image_encoder", torch_dtype=torch.float32) | |
| vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32) | |
| pipe = WanImageToVideoPipeline.from_pretrained( | |
| MODEL_ID, | |
| vae=vae, | |
| image_encoder=image_encoder, | |
| torch_dtype=torch.bfloat16 | |
| ) | |
| pipe.enable_model_cpu_offload() | |
| print("Pipeline loaded successfully.") | |
| # Increased duration to 120s just in case the free tier allows it | |
| def process_and_generate(image): | |
| if image is None: | |
| return None, None | |
| # Step 1: Isolate Product | |
| original_img = image.convert("RGB") | |
| cutout = remove(original_img) | |
| # Step 2: Prepare image | |
| img = cutout.convert("RGB") | |
| orig_w, orig_h = img.size | |
| aspect = orig_w / orig_h | |
| if aspect >= 1: | |
| new_w, new_h = 832, 480 | |
| else: | |
| new_w, new_h = 480, 832 | |
| img = img.resize((new_w, new_h), Image.LANCZOS) | |
| # Step 3: Generate Video | |
| try: | |
| prompt = "product rotating slowly 360 degrees on a turntable, studio lighting, seamless white background, smooth constant rotation, no morphing" | |
| print("Starting generation...") | |
| # REDUCED frames to 9 and steps to 20 to beat the 60-second timeout | |
| output = pipe( | |
| image=img, | |
| prompt=prompt, | |
| height=new_h, | |
| width=new_w, | |
| num_frames=9, | |
| num_inference_steps=20, | |
| guidance_scale=5.0 | |
| ).frames[0] | |
| # FIX: Save to the main folder, NOT a temp file | |
| video_path = "final_ad.mp4" | |
| export_to_video(output, video_path, fps=15) | |
| print("Generation complete! Video saved to:", video_path) | |
| return cutout, video_path | |
| except Exception as e: | |
| print(f"CRITICAL GENERATION ERROR: {e}") | |
| return cutout, None | |
| with gr.Blocks(title="AdSpin Studio") as demo: | |
| gr.Markdown("# 🛍️ AdSpin Studio") | |
| gr.Markdown("Upload a product photo. We isolate it, then Wan 2.1 animates it.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_img = gr.Image(type="pil", label="Upload Product Photo") | |
| generate_btn = gr.Button("Generate Video Ad", variant="primary") | |
| with gr.Column(): | |
| output_img = gr.Image(label="1. Product Isolated", type="pil") | |
| output_vid = gr.Video(label="2. Final Video Ad", format="mp4") | |
| generate_btn.click(fn=process_and_generate, inputs=input_img, outputs=[output_img, output_vid]) | |
| demo.launch() |