import gradio as gr
import torch
from diffusers import StableDiffusion3ControlNetPipeline, SD3ControlNetModel, UniPCMultistepScheduler
from huggingface_hub import login
import os
import spaces
from diffusers import StableDiffusion3ControlNetPipeline
from diffusers.models import SD3ControlNetModel, SD3MultiControlNetModel
from diffusers.utils import load_image

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

controlnet = SD3ControlNetModel.from_pretrained("InstantX/SD3-Controlnet-Tile")
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()