import gradio as gr
import torch
from diffusers import StableDiffusion3ControlNetPipeline, SD3ControlNetModel, UniPCMultistepScheduler
from huggingface_hub import login
import os
import spaces
from diffusers.utils import load_image, make_image_grid
import torch
from diffusers import StableDiffusionXLAdapterPipeline,T2IAdapter
from diffusers.models import T2IAdapter
from diffusers.schedulers import UniPCMultistepScheduler

# Log in to Hugging Face with your token
token = os.getenv("HF_TOKEN")
login(token=token)



# Load the T2I-Style Adapter and the SDXL pipeline
adapter = T2IAdapter.from_pretrained("TencentARC/t2i-adapter-style-sdxl")
pipe = StableDiffusionXLAdapterPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    adapter=adapter,
)

# Set up the scheduler and device
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.to("cuda", torch.float16)

# controlnet = SD3ControlNetModel.from_pretrained("alimama-creative/SD3-Controlnet-Softedge", torch_dtype=torch.float16)
#
# pipe = StableDiffusion3ControlNetPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", controlnet=controlnet)
# pipe.to("cuda", torch.float16)


@spaces.GPU
def generate_image(prompt, reference_image, controlnet_conditioning_scale):

    # Generate the image with ControlNet conditioning
    generated_image = pipe(
        prompt=prompt,
        control_image=load_image(reference_image),
        controlnet_conditioning_scale=controlnet_conditioning_scale,
    ).images[0]
    return generated_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.Slider(label="Control Net Conditioning Scale", minimum=0, maximum=1.0, step=0.1, value=0.6),
    ],
    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()