import torch
from diffusers import DiffusionPipeline, ControlNetModel
from diffusers.utils import load_image
import gradio as gr


controlnet = ControlNetModel.from_pretrained(
    "jasperai/Flux.1-dev-Controlnet-Upscaler",
    torch_dtype=torch.bfloat16
)

pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    controlnet=controlnet,
    torch_dtype=torch.bfloat16
)
pipe.to("cpu")

def enhance_image(input_image):
    control_image = load_image(input_image)

    w, h = control_image.size
    control_image = control_image.resize((w * 4, h * 4))

    image = pipe(
        prompt="", 
        control_image=control_image,
        controlnet_conditioning_scale=0.6,
        num_inference_steps=28, 
        guidance_scale=3.5,
        height=control_image.size[1],
        width=control_image.size[0]
    ).images[0]

    return image

interface = gr.Interface(
    fn=enhance_image,
    inputs=gr.inputs.Image(type="filepath"),
    outputs=gr.outputs.Image(type="pil"),
    title="Image Enhancer using FLUX.1-dev",
    description="Upload an image to enhance using the FLUX.1-dev model."
)

interface.launch()