Create Paintify.py

Paintify.py

@@ -0,0 +1,26 @@
+import torch
+from diffusers import StableDiffusionImg2ImgPipeline, DPMSolverMultistepScheduler
+from PIL import Image
+import os
+
+MODEL_ID = "runwayml/stable-diffusion-v1-5"
+LORA_PATH = "./models/ms_paint_style.safetensors"
+
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    safety_checker=None,
+)
+pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load LoRA
+pipe.load_lora_weights(LORA_PATH)
+pipe.fuse_lora()
+
+def stylize_image(image: Image.Image) -> Image.Image:
+    image = image.convert("RGB").resize((512, 512))
+    prompt = "MS Paint drawing, crude lines, fanart, child-like, cartoonish, naive"
+    result = pipe(prompt=prompt, image=image, strength=0.75, guidance_scale=7.5).images[0]
+    return result
+