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


model_id = "CompVis/stable-diffusion-v1-4"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")

def stylize_image(input_image, prompt):
    input_image = input_image.convert("RGB")
    input_image = input_image.resize((512, 512))
    output = pipe(prompt=prompt, image=input_image, strength=0.75).images[0]
    return output


iface = gr.Interface(
    fn=stylize_image, 
    inputs=[
        gr.Image(type="pil", label="Upload your image"),  # English label for image upload
        gr.Textbox(placeholder="Enter the art style... (e.g., Van Gogh style)", label="Art Style", lines=1)  # English label and example prompt
    ],
    outputs=gr.Image(label="Stylized Image"),  # English label for output
    title="Art and Style Transfer Demo",  # English title
    description="This demo uses the Stable Diffusion model to transform an image into a specified art style. Upload an image and enter a style prompt to get started.",  
    examples=[
        ["ben-grayland-gD-TjgDW0so-unsplash.jpg", "Van Gogh style"],  
    ]
)

iface.launch(share=True)