import gradio as gr
import torch
from diffusers import StableDiffusion3ControlNetPipeline, SD3ControlNetModel, UniPCMultistepScheduler
from huggingface_hub import login
import os
import spaces
# Log in to Hugging Face with your token
token = os.getenv("HF_TOKEN")
login(token=token)

# Model IDs for Stable Diffusion 1.5 and ControlNet
model_id = "stabilityai/stable-diffusion-3.5-large-turbo"
controlnet_id = "InstantX/SD3-Controlnet-Tile"

# Load the ControlNet model and Stable Diffusion pipeline
controlnet = SD3ControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch.float16)
pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
    model_id, controlnet=controlnet, torch_dtype=torch.float16
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

@spaces.GPU
def generate_image(prompt, reference_image,controlnet_conditioning_scale):
    # Prepare the reference image for ControlNet
    reference_image = reference_image.convert("RGB").resize((1024, 1024))

    # Generate the image with ControlNet conditioning
    generated_image = pipe(
        prompt=prompt,
        control_image=reference_image,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        guidance_scale=7.5,
        num_inference_steps=50
    ).images[0]
    return generated_image

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

# Launch the Gradio interface
interface.launch()