Update app.py

app.py

@@ -1,33 +1,33 @@
 import gradio as gr
 import torch
-import numpy as np
-from PIL import Image
 from diffusers import StableDiffusionInpaintPipeline
+from PIL import Image
+import numpy as np
 import cv2
 import os
 
-# Load a small CPU-friendly inpainting model
-model_id = "runwayml/stable-diffusion-inpainting"
-pipe = StableDiffusionInpaintPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
+# 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):
-    # Resize input
+def generate_drawing_video(image: Image.Image):
+    # Resize to manageable size for CPU
     image = image.resize((256, 256))
     
-    # Create a blank white mask
+    # Blank mask for inpainting
     mask = Image.new("L", image.size, color=255)
     
-    video_frames = []
+    frames = []
     
-    # Split mask into blocks to simulate drawing
+    # 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)
-    img_array = np.array(image)
     
     for i in range(0, len(flat_indices), pixels_per_frame):
         indices = flat_indices[i:i+pixels_per_frame]
@@ -36,27 +36,28 @@ def generate_drawing_video(image):
             y = idx // width
             current_mask.putpixel((x, y), 255)
         
-        # Generate partial "drawing" with AI inpainting
+        # Generate partial drawing
         out_img = pipe(prompt="pencil sketch", image=image, mask_image=current_mask).images[0]
-        video_frames.append(cv2.cvtColor(np.array(out_img), cv2.COLOR_RGB2BGR))
+        frames.append(cv2.cvtColor(np.array(out_img), cv2.COLOR_RGB2BGR))
     
     # Save video
     out_path = "/tmp/output.avi"
     fourcc = cv2.VideoWriter_fourcc(*"XVID")
-    out = cv2.VideoWriter(out_path, fourcc, 15, (image.width, image.height))
-    for frame in video_frames:
-        out.write(frame)
-    out.release()
+    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 Simulator",
-    description="Upload an image and watch an AI simulate a pencil drawing video.",
+    title="AI Hand Drawing Video",
+    description="Upload an image and watch an AI simulate a pencil drawing video on it.",
 )
 
-iface.launch()
-
+if __name__ == "__main__":
+    iface.launch()