# app.py

import os
import random
import uuid
import base64
import numpy as np
from PIL import Image
import torch
import streamlit as st

from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler

DESCRIPTION = """# DALL•E 3 XL v2 High Fi"""

def create_download_link(filename):
    with open(filename, "rb") as file:
        encoded_string = base64.b64encode(file.read()).decode('utf-8')
        download_link = f'<a href="data:image/png;base64,{encoded_string}" download="{filename}">Download Image</a>'
        return download_link
        
def save_image(img, prompt):
    unique_name = str(uuid.uuid4()) + ".png"
    img.save(unique_name)

    # save with prompt to save prompt as image file name
    filename = f"{prompt}.png"
    img.save(filename)
    return filename
    
def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

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

if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"

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

USE_TORCH_COMPILE = 0
ENABLE_CPU_OFFLOAD = 0



if torch.cuda.is_available():
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "fluently/Fluently-XL-v4",
        torch_dtype=torch.float16,
        use_safetensors=True,
    )
    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
    
    
    pipe.load_lora_weights("ehristoforu/dalle-3-xl-v2", weight_name="dalle-3-xl-lora-v2.safetensors", adapter_name="dalle")
    pipe.set_adapters("dalle")

    pipe.to("cuda")
    
    def generate(
        prompt: str,
        negative_prompt: str = "",
        use_negative_prompt: bool = False,
        seed: int = 0,
        width: int = 1024,
        height: int = 1024,
        guidance_scale: float = 3,
        randomize_seed: bool = False,
    ):
        seed = int(randomize_seed_fn(seed, randomize_seed))

        if not use_negative_prompt:
            negative_prompt = ""  # type: ignore

        images = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            guidance_scale=guidance_scale,
            num_inference_steps=20,
            num_images_per_prompt=1,
            cross_attention_kwargs={"scale": 0.65},
            output_type="pil",
        ).images
        image_paths = [save_image(img, prompt) for img in images]

        download_links = [create_download_link(path) for path in image_paths]

        print(image_paths)
        return image_paths, seed, download_links
else:
    st.warning("CUDA is not available. The demo may not work on CPU.")

    

examples = [
"a modern hospital room with advanced medical equipment and a patient resting comfortably",
"a team of surgeons performing a delicate operation using state-of-the-art surgical robots",
"a elderly woman smiling while a nurse checks her vital signs using a holographic display",
"a child receiving a painless vaccination from a friendly robot nurse in a colorful pediatric clinic",
"a group of researchers working in a high-tech laboratory, developing new treatments for rare diseases",
"a telemedicine consultation between a doctor and a patient, using virtual reality technology for a immersive experience"
]

#st.set_page_config(page_title="DALL•E 3 XL v2 High Fi", layout="centered")
st.markdown(DESCRIPTION)

with st.form(key="generation_form"):
    prompt = st.text_input("Prompt", max_chars=None, placeholder="Enter your prompt")
    use_negative_prompt = st.checkbox("Use negative prompt", value=True)
    if use_negative_prompt:
        negative_prompt = st.text_area(
            "Negative prompt",
            value="""(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, (NSFW:1.25)""",
            placeholder="Enter a negative prompt",
            height=100,
        )
    else:
        negative_prompt = ""
    
    col1, col2 = st.columns(2)
    with col1:
        width = st.slider("Width", min_value=512, max_value=2048, step=8, value=1920)
    with col2:    
        height = st.slider("Height", min_value=512, max_value=2048, step=8, value=1080)
    
    col3, col4 = st.columns(2)
    with col3:
        guidance_scale = st.slider("Guidance Scale", min_value=0.1, max_value=20.0, step=0.1, value=20.0)
    with col4:
        randomize_seed = st.checkbox("Randomize seed", value=True)
        if not randomize_seed:
            seed = st.slider("Seed", min_value=0, max_value=MAX_SEED, step=1, value=0)
        else:
            seed = 0
    
    run_button = st.form_submit_button("Run")

if run_button:
    image_paths, seed, download_links = generate(
        prompt=prompt,
        negative_prompt=negative_prompt,
        use_negative_prompt=use_negative_prompt,
        seed=seed,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        randomize_seed=randomize_seed,
    )
    
    for image_path in image_paths:
        st.image(image_path, caption=prompt)
    
    for download_link in download_links:
        st.markdown(download_link, unsafe_allow_html=True)
    
    st.text(f"Seed: {seed}")
    

def set_prompt(example):
    prompt.value = example

st.subheader("Examples")
for example in examples:
    st.button(example, key=example, on_click=set_prompt, args=(example,))