import torch
import gradio as gr
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from PIL import Image

# Load ControlNet model
controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16)

# Load Stable Diffusion Pipeline
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    controlnet=controlnet,
    torch_dtype=torch.float16
).to("cuda")

# Function to process images
def remove_filter(image, prompt):
    output = pipe(prompt=prompt, image=image).images[0]
    return output

# Gradio Interface
iface = gr.Interface(
    fn=remove_filter,
    inputs=[
        gr.Image(type="pil"),  # Image input
        gr.Textbox(label="Custom Prompt", placeholder="Enter a custom prompt", value="restore the original colors and details, remove the filter")  # Textbox for the prompt
    ],
    outputs=gr.Image(type="pil"),
    title="Filter Removal using ControlNet",
    description="Upload a filtered image and provide a prompt to restore the original image.",
)

# Launch the app
if __name__ == "__main__":
    iface.launch()