import random

import gradio as gr
import spaces
import torch
from diffusers import Flux2KleinPipeline
import os

MODEL_ID = "alfredplpl/ecocoro"

DEFAULT_PROMPT = "recent, highreso, 1girl, solo, close-up, brown hair, blue eyes, white T-shirts, ahoge, looking at viewer, blue background, simple background"
ASPECT_RATIOS = {
    "Square": (1024, 1024),
    "Portrait": (768, 1344),
    "Landscape": (1344, 768),
}

token=os.environ["HF_TOKEN"]

pipe = Flux2KleinPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
).to("cuda")


def get_duration(prompt, aspect_ratio):
    return 120


@spaces.GPU(duration=get_duration)
def generate(prompt, aspect_ratio):
    width, height = ASPECT_RATIOS[aspect_ratio]
    
    seed = random.randint(0, 2**32 - 1)
    generator = torch.Generator(device="cuda").manual_seed(seed)

    negative_prompt="bad anatomy, low quality, bad hands, extra fingers"
    
    image = pipe(
        prompt=prompt,
        height=height,
        width=width,
        num_inference_steps=50,
        guidance_scale=4.0,
        generator=generator,
    ).images[0]

    return image, seed


with gr.Blocks(title="Ecocoro") as demo:
    gr.Markdown(
        """
# Ecocoro

A simple text-to-image demo for [alfredplpl/ecocoro-preview-1](https://huggingface.co/alfredplpl/ecocoro-preview-1).
"""
    )

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                value=DEFAULT_PROMPT,
                lines=6,
                placeholder="Enter your prompt here...",
            )

            aspect_ratio = gr.Radio(
                label="Aspect Ratio",
                choices=["Square", "Portrait", "Landscape"],
                value="Square",
            )

            generate_button = gr.Button("Generate", variant="primary")

            seed = gr.Number(
                label="Seed",
                precision=0,
                visible=False,
            )

        with gr.Column(scale=1):
            output_image = gr.Image(
                label="Generated Image",
                type="pil",
            )

    gr.Examples(
        examples=[
            [
                "recent, highreso, 1girl, solo, close-up, brown hair, blue eyes, white T-shirts, ahoge, blue background, simple background",
                "Square",
            ],
            [
                "recent, highreso, 1girl, solo, upper body, black hair, red eyes, school uniform, soft lighting, simple background",
                "Portrait",
            ],
            [
                "recent, highreso, 1girl, solo, sitting by the window, sunset, warm lighting, anime style",
                "Landscape",
            ],
        ],
        inputs=[prompt, aspect_ratio],
    )

    generate_button.click(
        fn=generate,
        inputs=[prompt, aspect_ratio],
        outputs=[output_image, seed],
    )


demo.queue(max_size=10).launch()