app.py

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


model_id = "CompVis/stable-diffusion-v1-4"


if torch.cuda.is_available():
    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
    pipe = pipe.to("cuda")
else:
    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id)
    pipe = pipe.to("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"),
        gr.Textbox(placeholder="Enter the art style... (e.g., Van Gogh style)", label="Art Style", lines=1)
    ],
    outputs=gr.Image(label="Stylized Image"),
    title="Art and Style Transfer Demo",
    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)