import gradio as gr
import gradio.components as gc
import torch
import numpy as np
from diffusers import DiffusionPipeline
from huggingface_hub import login, HfApi, HfFolder
from PIL import Image
import os
from datetime import datetime
import shutil

# Get your Hugging Face API token
folder = HfFolder()
token = folder.get_token()

# Instantiate the Hugging Face API
api = HfApi()

login(token=os.environ.get('HF_KEY'))

device = "cuda" if torch.cuda.is_available() else "cpu"

torch.cuda.max_memory_allocated(device=device)

pipe1 = DiffusionPipeline.from_pretrained("FFusion/FFusionXL-BASE", torch_dtype=torch.float16, use_safetensors=True)
pipe2 = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)

pipe1 = pipe1.to(device)
pipe1.enable_xformers_memory_efficient_attention()

pipe2 = pipe2.to(device)
pipe2.enable_xformers_memory_efficient_attention()

def save_image_to_hf_space(image_np, image_name):
    # Name of your Hugging Face repo
    repo_name = "FFusion/FF2"

    # Convert the numpy array to an image
    image = Image.fromarray(image_np.astype('uint8'))

    # Append a timestamp to the filename
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    image_name_with_timestamp = f"{image_name}-{timestamp}.png"

    # Save the image locally
    local_image_path = f"./{image_name_with_timestamp}"
    image.save(local_image_path)

    # Upload the image to your Hugging Face repo
    api.upload_file(
        token=token,
        path_or_fileobj=local_image_path,
        repo_id=repo_name,
        path_in_repo=image_name_with_timestamp  # The path where the image will be stored in the repository
    )

    # Save the image to the persistent storage
    persistent_storage_path = f"/data/{image_name_with_timestamp}"
    shutil.copy(local_image_path, persistent_storage_path)

   
article = f"""
<div style="text-align: center;">
    <img src="https://cdn-uploads.huggingface.co/production/uploads/6380cf05f496d57325c12194/LIONhgnyxUuEynNivGsME.jpeg" alt="ffusionAI-FFusionXL-SDXL-preview" width="100%" height="600">
</div>
<div style="display: flex; flex-direction: column; align-items: center; gap: 10px;">
    <h2>Citation</h2>
    <p>Please note that the demo is intended solely for academic and research purposes. This demo features the FFusionXL-BASE model developed by FFusion.AI, a division of Source Code Bulgaria Ltd.</p>
    <h2>Acknowledgement of Original Work and Modifications</h2>
    <p>This Software is based on the Stable Diffusion XL Base 1.0 model developed by Stability AI Ltd. FFusion AI and Source Code Bulgaria Ltd. have made modifications and enhancements to the original Software for the creation of the FFusionXL-BASE model...</p>
    <h2>Warning and Compliance</h2>
    <p><span style="color: red; font-weight: bold;">All images and content generated through this demo are logged in a Hugging Face repository, and we actively monitor for violations of these terms.</span> Any use of the demo for generating inappropriate or unlawful content is strictly prohibited...</p>
    <h2>License</h2>
    <p><a href="https://huggingface.co/stabilityai/stable-diffusion-xl-base-0.9/blob/main/LICENSE.md">SDXL 0.9 Research License</a></p>
    <p><a href="https://huggingface.co/FFusion/FFusionXL-09-SDXL/blob/main/LICENSE.md">FFXL 0.9 Research License</a></p>
    <div style="display: flex; flex-wrap: wrap; gap: 2px;">
        <img src="https://img.shields.io/badge/🔥 Refiner Compatible-Yes-success">
        <img src="https://img.shields.io/badge/💻 CLIP--ViT%2FG and CLIP--ViT%2FL tested-Yes-success">
        <img src="https://img.shields.io/badge/🧨 FFXL Diffusers-available-brightgreen">
    </div>
    <div style="display: flex; flex-wrap: wrap; gap: 2px;">
        <a href="https://github.com/1e-2" target="_new" rel="ugc"><img src="https://img.shields.io/badge/GitHub-1e--2-green"></a>
        <a href="https://www.facebook.com/FFusionAI/" target="_new" rel="ugc"><img src="https://img.shields.io/badge/Facebook-FFusionAI-blue"></a>
        <a href="https://civitai.com/models/82039/ffusion-ai-sd-21" target="_new" rel="ugc"><img src="https://img.shields.io/badge/Civitai-FFusionAI-blue"></a>
    </div>
    <a href="mailto:di@ffusion.ai"><img src="https://img.shields.io/badge/Email-di%40ffusion.ai-blue?style=for-the-badge&logo=gmail"></a>
</div>

"""

# Create the Gradio interface
gr.Interface(fn=genie, 
             inputs=[gr.Textbox(label='Describe your FFusion idea. 77 Token Limit.', lines=2), 
                     gr.Textbox(label='Things the AI should not create (negative prompt)', lines=2), 
                     gr.Slider(1, 15, 10), gr.Slider(25, maximum=100, value=50, step=1), 
                     gr.Slider(minimum=1, step=1, maximum=999999999999999999, randomize=True)], 
             outputs=[gc.Image(type='numpy', label="FFusionXL Base Image"), gc.Image(type='numpy', label="Refined Image")],
             title="FFusionXL Base - Generate and Refine", 
             description='<div style="display: flex; flex-wrap: wrap; gap: 2px; justify-content: center;"><a href="https://huggingface.co/FFusion/FFusionXL-BASE" target="_new" rel="ugc"><img src="https://img.shields.io/badge/FFusionXL--BASE--SDXL-Model-pink" alt="FFusionXL-BASE-SDXL"></a>  <a href="https://huggingface.co/FFusion/FFusionXL-09-SDXL/blob/main/LICENSE.md" target="_new" rel="ugc"><img src="https://img.shields.io/badge/License-FFXL%20Research%20License-blue"></a></div>', 
             article=article, 
             css="""
                 .gr-textbox {
                     width: 100%;
                 }
                 .gr-image {
                     max-width: 50%;
                     display: inline-block;
                 }
             """
).launch(debug=True, max_threads=10)