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
@spaces.GPU(duration=120) 
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()