import torch
import spaces
import os
from diffusers.utils import load_image
from diffusers import FluxControlNetModel, FluxControlNetPipeline
import gradio as gr
huggingface_token = os.getenv("HUGGINFACE_TOKEN")

# Load pipeline
controlnet = FluxControlNetModel.from_pretrained(
    "jasperai/Flux.1-dev-Controlnet-Upscaler",
    torch_dtype=torch.bfloat16
)
pipe = FluxControlNetPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    controlnet=controlnet,
    torch_dtype=torch.bfloat16,
    token=huggingface_token
)
pipe.to("cuda")

@spaces.GPU
def generate_image(prompt, control_image):
    # Load control image
    control_image = load_image(control_image)
    w, h = control_image.size
    # Upscale x1
    control_image = control_image.resize((w * 1, h * 1))
    image = pipe(
        prompt=prompt,
        control_image=control_image,
        controlnet_conditioning_scale=0.6,
        num_inference_steps=8,
        guidance_scale=3.5,
        height=control_image.size[1],
        width=control_image.size[0]
    ).images[0]
    return image

# Create Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
        gr.Image(type="pil", label="Control Image"),
    ],
    outputs=[
        gr.Image(type="pil", label="Generated Image"),
    ],
    title="FLUX ControlNet Image Generation",
    description="Generate images using the FluxControlNetPipeline. Upload a control image and enter a prompt to create an image.",
)

# Launch the app
iface.launch()