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