sphinxsolution's picture
updated app.py
1a2fd9e verified
Raw History Blame Contribute Delete
1.15 kB
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()