import gradio as gr
import numpy as np
import spaces
import torch
import random
from diffusers import FluxFillPipeline
from PIL import Image
import os

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

# Debugging CUDA environment and decorator
print("CUDA_VISIBLE_DEVICES:", os.environ.get("CUDA_VISIBLE_DEVICES", "Not set"))
print("CUDA Available:", torch.cuda.is_available())
print("GPU Count:", torch.cuda.device_count())
print(
    "Current Device:",
    torch.cuda.current_device() if torch.cuda.is_available() else "None",
)
print("Spaces GPU Decorator Active: True")


def calculate_optimal_dimensions(image: Image.Image):
    original_width, original_height = image.size
    MIN_ASPECT_RATIO = 9 / 16
    MAX_ASPECT_RATIO = 16 / 9
    FIXED_DIMENSION = 1024
    original_aspect_ratio = original_width / original_height
    if original_aspect_ratio > 1:  # Wider than tall
        width = FIXED_DIMENSION
        height = round(FIXED_DIMENSION / original_aspect_ratio)
    else:  # Taller than wide
        height = FIXED_DIMENSION
        width = round(FIXED_DIMENSION * original_aspect_ratio)
    width = (width // 8) * 8
    height = (height // 8) * 8
    calculated_aspect_ratio = width / height
    if calculated_aspect_ratio > MAX_ASPECT_RATIO:
        width = (height * MAX_ASPECT_RATIO // 8) * 8
    elif calculated_aspect_ratio < MIN_ASPECT_RATIO:
        height = (width / MIN_ASPECT_RATIO // 8) * 8
    width = max(width, 576) if width == FIXED_DIMENSION else width
    height = max(height, 576) if height == FIXED_DIMENSION else height
    return width, height


def create_full_mask(image):
    """Generate a fully white mask for the entire image."""
    return Image.fromarray(
        np.ones((image.height, image.width), dtype=np.uint8) * 255
    ).convert("L")


@spaces.GPU  # Removed duration parameter
def infer(
    image,
    prompt,
    seed=42,
    randomize_seed=False,
    width=1024,
    height=1024,
    guidance_scale=50,
    num_inference_steps=28,
):
    if not image:
        raise gr.Error("Please upload an image.")
    mask = create_full_mask(image)  # Auto-generate full white mask
    width, height = calculate_optimal_dimensions(image)
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    pipe = FluxFillPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-FILL-dev", torch_dtype=torch.bfloat16
    ).to("cuda")
    generator = torch.Generator("cuda").manual_seed(seed)
    try:
        image = pipe(
            prompt=prompt,
            image=image,
            mask_image=mask,
            height=height,
            width=width,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            max_sequence_length=512,
            generator=generator,
        ).images[0]
        return image, seed
    except Exception as e:
        raise gr.Error(f"Error during inpainting: {str(e)}")


examples = [
    "Add a glowing crescent moon on the forehead, vivid red eyes, and a shadowy, dark, misty background to the subject, while preserving the exact structure and details of the original human face, animal, or object",
    "a tiny astronaut hatching from an egg on the moon",
    "a cat holding a sign that says hello world",
    "an anime illustration of a wiener schnitzel",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 1000px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            f"""# FLUX.1 Fill [dev]
12B param rectified flow transformer structural conditioning tuned, guidance-distilled from [FLUX.1 [pro]](https://huggingface.co/black-forest-labs/FLUX.1-pro)
[[non-commercial license](https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev/blob/main/LICENSE.txt)] [[blog](https://bfl.ai/blog/2024/08/01/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev)]
"""
        )
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(
                    label="Upload image for editing",
                    type="pil",
                    sources=["upload"],
                    image_mode="RGB",
                    height=600,
                )
                prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter your prompt",
                    container=False,
                )
                gr.Examples(
                    examples=examples,
                    inputs=[prompt],
                    label="Example Prompts",
                )
                run_button = gr.Button("Run")
                result = gr.Image(label="Result", show_label=False)
        with gr.Accordion("Advanced Settings", open=False):
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                    visible=False,
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                    visible=False,
                )
            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
                    maximum=50,
                    step=0.5,
                    value=50,
                )
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )
        gr.on(
            triggers=[run_button.click, prompt.submit],
            fn=infer,
            inputs=[
                image_input,
                prompt,
                seed,
                randomize_seed,
                width,
                height,
                guidance_scale,
                num_inference_steps,
            ],
            outputs=[result, seed],
        )

demo.launch()