import gradio as gr
from huggingface_hub import login
import os
import spaces
from diffusers import AutoPipelineForText2Image
from diffusers.utils import load_image
import torch
import tempfile

token = os.getenv("HF_TOKEN")
login(token=token)


pipeline = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16).to("cuda")
pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin")



@spaces.GPU
def generate_image(prompt, reference_image, controlnet_conditioning_scale):
    style_images = [load_image(f.name) for f in reference_image]

    pipeline.set_ip_adapter_scale(controlnet_conditioning_scale)

    image = pipeline(
        prompt=prompt,
        ip_adapter_image=[style_images],
        negative_prompt="",
        guidance_scale=5,
        num_inference_steps=30,
    ).images[0]

    return image

# Set up Gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt"),
        # gr.Image( type= "filepath",label="Reference Image (Style)"),
        gr.File(file_count="multiple",label="Reference Image (Style)"),
        gr.Slider(label="Control Net Conditioning Scale", minimum=0, maximum=1.0, step=0.1, value=1.0),
    ],
    outputs="image",
    title="Image Generation with Stable Diffusion 3 medium and ControlNet",
    description="Generates an image based on a text prompt and a reference image using Stable Diffusion 3 medium with ControlNet."

)

interface.launch()