import torch
import gradio as gr
from diffusers import DiffusionPipeline
import os
import spaces
import accelerate

# UI
DESCRIPTION = '''
<div>
<h1 style="text-align: center;">Osiris 🌌</h1>
<p>This has an open source stable diffuser from <a href="https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0"><b>stable-diffusion-xl-base-1.0</b></a></p>
</div>
'''

# load both base and refiner
base = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16").to('cuda')
refiner = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0",
                                            text_encoder_2=base.text_encoder_2,
                                            vae=base.vae,
                                            torch_dtype=torch.float16,
                                            use_safetensor=True,
                                            variant="fp16").to('cuda')

# function to take input and generate text tokena
@spaces.GPU(duration=120)
def osiris(prompt: str,
           n_steps: int,
           high_noise_frac: float):
    """
    Takes input, passes it into the pipeline, 
    get the top 5 scores, and ouput those scores into images
    """

    # n_steps = int(n_steps)

    # Generate image based on text
    image_base = base(
        prompt=prompt,
        num_inference_steps=n_steps,
        denoising_end=high_noise_frac,
        output_type="latent"
    ).images
    image = refiner(
        prompt=prompt,
        num_inference_steps=n_steps,
        denoising_start=high_noise_frac,
        image=image_base
    ).images[0]

    return image

with gr.Blocks(fill_height=True) as demo:
    gr.Markdown(DESCRIPTION)
    gr.Interface(
        fn=osiris,
        inputs="text",
        outputs="image",
        fill_height=True,
        additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
        additional_inputs=[
            gr.Slider(minimum=20,
                      maximum=100,
                      step=1,
                      value=40,
                      label="Number of Inference Steps",
                      render=False),
            gr.Slider(minimum=0.0,
                      maximum=1.0,
                      step=0.1,
                      value=0.8,
                      label="High Noise Fraction",
                      render=False),
        ],
        examples=[
            ["A sprawling cyberpunk metropolis at dusk, with towering skyscrapers adorned with neon signs, holographic billboards, and flying cars weaving through the sky. On a crowded street corner, a cybernetically enhanced street artist creates mesmerizing light sculptures with their augmented reality gloves. Rain glistens on the bustling sidewalks as pedestrians with colorful umbrellas rush past."],
            ["A mystical enchanted forest bathed in twilight, where bioluminescent plants cast an ethereal glow. A crystal-clear waterfall cascades into a shimmering pool, surrounded by ancient trees with twisted roots. A lone elf archer, dressed in elegant green and gold attire, stands on a moss-covered rock, her bow drawn as she watches over the tranquil scene. Ethereal fairies with delicate wings flutter around, leaving trails of sparkling dust in the air."],
            ["An elaborate steampunk dirigible floating gracefully above a cloud-covered landscape. The airship, with its brass and copper gears, massive steam-powered engines, and ornate Victorian design, cruises past a golden sunset. A distinguished gentleman in a top hat and goggles stands on the observation deck, holding a brass spyglass to his eye as he surveys the horizon. Passengers in vintage attire marvel at the view."],
            ["A charming, snow-covered log cabin nestled in the heart of a tranquil mountain range during a starry winter night. Smoke gently curls from the stone chimney, and warm light spills from the windows, illuminating the cozy interior. A young woman in a thick woolen coat and a fur-lined hat stands by the front door, holding a lantern that casts a warm glow on the snow. Pine trees, heavy with snow, frame the scene, while the Northern Lights dance across the sky."],
            ["A vibrant underwater kingdom where a colorful coral reef teems with marine life. Schools of iridescent fish swim through the crystal-clear waters, and a sunken pirate ship, encrusted with barnacles and treasure chests, rests on the sandy seabed. A curious mermaid with flowing turquoise hair and a shimmering silver tail explores the depths, holding a glowing pearl in her hand. Playful dolphins swim around her, and a wise old sea turtle watches from a nearby rock."]
        ],
        cache_examples=False
    )

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