import spaces
import gradio as gr
import numpy as np
from diffusers import AutoPipelineForText2Image
import torch
from huggingface_hub import snapshot_download
from io import BytesIO
from PIL import Image

SPACE_USERNAME = 'KR_4dmin'
SPACE_PASSWORD = 'KR_4dmin'

snapshot_download(repo_id="Roomie/xavyy", cache_dir='./')

pipeline = AutoPipelineForText2Image.from_pretrained(
    'black-forest-labs/FLUX.1-schnell', torch_dtype=torch.bfloat16).to('cuda')
pipeline.load_lora_weights('Roomie/xavyy', weight_name='xavyy.safetensors')


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


@spaces.GPU
def infer(prompt, height, width, guidance_scale, num_inference_steps):

    image = pipeline(
        prompt=prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        # image=refer_image
        # generator=generator
    ).images[0]

    return image


examples = [
    "Xavy, a virtual content creator, is in a high-tech futuristic studio filled with holographic screens and cutting-edge gadgets. He’s presenting the latest smartphone technology, wearing a sleek tech-inspired outfit with neon accents. The background features floating data, robots assisting him, and advanced digital tools. His facial expression is enthusiastic as he explains the potential of artificial intelligence in smartphones. The atmosphere is dynamic and full of futuristic energy. Negative prompt: avoid multiple versions of Xavy, avoid distorted facial features, malformed hands, excessive or broken gadgets, unrealistic proportions in the body or technology, extra limbs.",
    "Xavy stands on a stage at a technology innovation conference, speaking passionately about the future of AI in smartphones. Behind him, a massive screen displays 3D holographic models of a cutting-edge phone design. The audience is captivated as he gestures towards the hologram, explaining how AI will revolutionize user interaction. He’s wearing a sleek black outfit with a futuristic smartwatch, and the lighting is focused on him while the background is filled with technological elements like drones and digital billboards.  Negative prompt: avoid duplicated Xavy figures, warped or incomplete body parts, malformed facial expressions, extra gadgets or overlapping elements, unnatural lighting, broken equipment, unrealistic audience features."
]

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

with gr.Blocks(css=css) as demo:

    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""
        # Draw Virtual Creators
        """)

        with gr.Row():

            prompt = gr.TextArea(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )

        run_button = gr.Button("Run", scale=0)

        result = gr.Image(label="Result", show_label=True)

        with gr.Accordion("Advanced Settings", open=True):

            # seed = gr.Slider(
            #     label="Seed",
            #     minimum=0,
            #     maximum=MAX_SEED,
            #     step=1,
            #     value=0,
            #     visible=False
            # )

            # randomize_seed = gr.Checkbox(
            #     label="Randomize seed", value=True, visible=False)

            with gr.Row():

                width = gr.Slider(
                    label="Width",
                    info="Solo cambiar con las flechas o escoger una medida que sea multiplo de 8",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=8,
                    value=1024,
                )

                height = gr.Slider(
                    label="Height",
                    info="Solo cambiar con las flechas o escoger una medida que sea multiplo de 8",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=8,
                    value=1024,
                )

            with gr.Row():

                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    info="Valores mas altos se apega mas al prompt, la calidad del resultado baja. Valores bajos permite creatividad pero se aleja del prompt",
                    minimum=0.0,
                    maximum=10.0,
                    step=0.1,
                    value=3.5,
                    visible=True
                )

                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    info="Entre mas numeros de inferencia mejor calidad de la imagen. Toma mas tiempo generar la imagen.",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=48,
                    visible=True
                )

        gr.Examples(
            examples=examples,
            inputs=[prompt]
        )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[prompt, height, width, guidance_scale, num_inference_steps],
        outputs=[result]
    )

demo.queue().launch(share=True, auth=(SPACE_USERNAME, SPACE_PASSWORD))