import gradio as gr
import torch
from PIL import Image
from diffusers import StableDiffusionImg2ImgPipeline

# Load your model (replace with your own model repo)
model_id = "brothelsnsprout/Training_image_generator"

pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16
).to("cuda")

# Stylize function
def stylize_image(image: Image.Image, prompt: str):
    image = image.convert("RGB").resize((512, 512))
    result = pipe(prompt=prompt, image=image, strength=0.8, guidance_scale=7.5)
    return result.images[0]

# Gradio Interface
interface = gr.Interface(
    fn=stylize_image,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Textbox(label="Text Prompt (Style Instruction)", placeholder="e.g. in the style of Studio Ghibli")
    ],
    outputs=gr.Image(label="Styled Output"),
    title="Image-to-Image Styler",
    description="Upload an image and describe how it should be transformed in style.",
)

if __name__ == "__main__":
    interface.launch()