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 = "restore the original colors and details, remove the filter"
    output = pipe(prompt, image=image).images[0]
    return output

# Gradio Interface
iface = gr.Interface(
    fn=remove_filter,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Filter Removal using ControlNet",
    description="Upload a filtered image, and this app will attempt to restore the original image.",
)

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