tools/edit.py

import torch
from PIL import Image
from diffusers import AutoPipelineForImage2Image

pipe = None


def load_model():
    global pipe

    if pipe is None:
        pipe = AutoPipelineForImage2Image.from_pretrained(
            "stabilityai/sd-turbo",
            torch_dtype=torch.float32
        )
        pipe.to("cpu")

    return pipe


def edit_image(image_path: str, prompt: str):

    if image_path is None:
        return None

    model = load_model()

    image = Image.open(image_path).convert("RGB")

    result = model(
        prompt=prompt,
        image=image,
        strength=0.75,
        guidance_scale=3.0,
        num_inference_steps=20
    )

    return result.images[0]