app.py

import gradio as gr
import torch
from diffusers import StableDiffusionInpaintPipeline
from PIL import Image
import numpy as np
import cv2
import os

# Load a CPU-friendly Stable Diffusion inpainting model
pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float32
)
pipe.to("cpu")

def generate_drawing_video(image: Image.Image):
    # Resize to manageable size for CPU
    image = image.resize((256, 256))
    
    # Blank mask for inpainting
    mask = Image.new("L", image.size, color=255)
    
    frames = []
    
    # Create random pixel blocks to simulate hand drawing
    height, width = image.size
    pixels_per_frame = 500
    flat_indices = np.arange(height*width)
    np.random.shuffle(flat_indices)
    
    current_mask = Image.new("L", image.size, color=0)
    
    for i in range(0, len(flat_indices), pixels_per_frame):
        indices = flat_indices[i:i+pixels_per_frame]
        for idx in indices:
            x = idx % width
            y = idx // width
            current_mask.putpixel((x, y), 255)
        
        # Generate partial drawing
        out_img = pipe(prompt="pencil sketch", image=image, mask_image=current_mask).images[0]
        frames.append(cv2.cvtColor(np.array(out_img), cv2.COLOR_RGB2BGR))
    
    # Save video
    out_path = "/tmp/output.avi"
    fourcc = cv2.VideoWriter_fourcc(*"XVID")
    video_writer = cv2.VideoWriter(out_path, fourcc, 15, (image.width, image.height))
    for frame in frames:
        video_writer.write(frame)
    video_writer.release()
    
    return out_path

# Create Gradio interface
iface = gr.Interface(
    fn=generate_drawing_video,
    inputs=gr.Image(type="pil"),
    outputs=gr.Video(type="file"),
    title="AI Hand Drawing Video",
    description="Upload an image and watch an AI simulate a pencil drawing video on it.",
)

if __name__ == "__main__":
    iface.launch()