import gradio as gr
from diffusers import StableDiffusionImg2ImgPipeline
import torch
import base64

model = StableDiffusionImg2ImgPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.float16
).to("cuda" if torch.cuda.is_available() else "cpu")

def transform_image(image, prompt="place it"):
    result = model(prompt=prompt, image=image).images[0]
    return result

gr.Interface(
    fn=transform_image,
    inputs=[gr.Image(type="pil"), gr.Textbox(label="Prompt", value="place it")],
    outputs=gr.Image(type="pil", label="Modified Image"),
    title="Image Overlay with LoRA"
).launch()