import gradio as gr
from diffusers import AutoPipelineForImage2Image
import torch

pipe = AutoPipelineForImage2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16
).to("cpu")
#For now, we have to use cpu because CUDA is not available. This taks 8mins per image btw.
#pipe.load_lora_weights("enhanceaiteam/Flux-Uncensored-V2")

def process(image, prompt, strength):
    result = pipe(prompt=prompt, image=image, strength=strength)
    return result.images[0]

demo = gr.Interface(
    fn=process,
    inputs=[
        gr.Image(type="pil"),
        gr.Textbox(label="Prompt"),
        gr.Slider(0.1, 1.0, value=0.75, label="Strength")
    ],
    outputs=gr.Image()
)

demo.launch()