import os
import random

import spaces
import gradio as gr
import torch
from diffusers.utils import load_image
from diffusers.pipelines.flux.pipeline_flux_controlnet import FluxControlNetPipeline
from diffusers.models.controlnet_flux import FluxControlNetModel
import numpy as np
from huggingface_hub import login, snapshot_download


# Configuration
BASE_MODEL = 'black-forest-labs/FLUX.1-dev'
CONTROLNET_MODEL = 'promeai/FLUX.1-controlnet-lineart-promeai'
CSS = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

# Setup
AUTH_TOKEN = os.getenv("HF_AUTH_TOKEN")
if AUTH_TOKEN:
    login(AUTH_TOKEN)
else:
    raise ValueError("Hugging Face auth token not found. Please set HF_AUTH_TOKEN in the environment.")

MODEL_DIR = snapshot_download(
    repo_id=BASE_MODEL,
    revision="main",
    use_auth_token=AUTH_TOKEN
)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
TORCH_DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32

CONTROLNET = FluxControlNetModel.from_pretrained(CONTROLNET_MODEL, torch_dtype=TORCH_DTYPE)
PIPE = FluxControlNetPipeline.from_pretrained(MODEL_DIR, controlnet=CONTROLNET, torch_dtype=TORCH_DTYPE)
torch.cuda.empty_cache()
PIPE = PIPE.to(DEVICE)

MAX_SEED = np.iinfo(np.int32).max

@spaces.GPU(duration=140)
def infer(
    prompt,
    control_image_path,
    controlnet_conditioning_scale,
    guidance_scale,
    num_inference_steps,
    seed,
    randomize_seed,
):
    global DEVICE, TORCH_DTYPE
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    TORCH_DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
    print(f"Inference: using device: {DEVICE} (torch_dtype={TORCH_DTYPE})")
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.manual_seed(seed)
    control_image = load_image(control_image_path) if control_image_path else None

    # Generate image
    result = PIPE(
        prompt=prompt,
        control_image=control_image,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=generator,
    ).images[0]

    return result, seed

with gr.Blocks(css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# Flux.1[dev] LineArt")
        gr.Markdown("### Zero-shot Partial Style Transfer for Line Art Images, Powered by FLUX.1")
        control_image = gr.Image(
                sources=['upload', 'webcam', 'clipboard'],
                type="filepath",
                label="Control Image (LineArt)"
        )
        prompt = gr.Text(
            label="Prompt",
            placeholder="Enter your prompt",
            max_lines=1,
            container=False
        )
        run_button = gr.Button("Generate", variant="primary")
        result = gr.Image(label="Result", show_label=False)
        with gr.Accordion("Advanced Settings", open=False):
            controlnet_conditioning_scale = gr.Slider(
                label="ControlNet Conditioning Scale",
                minimum=0.0,
                maximum=1.0,
                value=0.6,
                step=0.1
            )
            guidance_scale = gr.Slider(
                label="Guidance Scale",
                minimum=1.0,
                maximum=10.0,
                value=3.5,
                step=0.1
            )
            num_inference_steps = gr.Slider(
                label="Number of Inference Steps",
                minimum=1,
                maximum=100,
                value=28,
                step=1
            )
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0
            )
            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)

        gr.Examples(
            examples=[
                "Shiba Inu wearing dinosaur costume riding skateboard",
                "Victorian style mansion interior with candlelight",
                "Loading screen for Grand Theft Otter: Clam Andreas"
            ],
            inputs=[prompt]
        )
        
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn = infer,
        inputs=[
            prompt,
            control_image,
            controlnet_conditioning_scale,
            guidance_scale,
            num_inference_steps,
            seed,
            randomize_seed
        ],
        outputs = [result, seed]
    )

if __name__ == "__main__":
    demo.launch()