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 add_watermark(image_np):
    img = Image.fromarray(image_np.astype('uint8'))
    draw = ImageDraw.Draw(img)
    watermark_text_line1 = "WARNING: This image is generated for Research & Demonstration Purposes Only."
    watermark_text_line2 = "Any misuse or inappropriate use and may be subject to legal action."
    font = ImageFont.truetype("arial.ttf", size=12)
    position_line1 = (10, img.height - 80)
    position_line2 = (10, img.height - 60)  # Adjust this value based on the font size and desired spacing
    color_line1 = "white"
    color_line2 = "black"
    draw.text(position_line1, watermark_text_line1, font=font, fill=color_line1)
    draw.text(position_line2, watermark_text_line2, font=font, fill=color_line2)
    return np.array(img)



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)

   


def genie(prompt, negative_prompt, prompt_2, negative_prompt_2, scale, guidance_scale, height, width, aesthetic_score, negative_aesthetic_score, steps, seed):
    torch.cuda.empty_cache()
    generator = torch.Generator(device=device).manual_seed(seed)
    int_images = pipe1(prompt=prompt, prompt_2=prompt_2, negative_prompt=negative_prompt, negative_prompt_2=negative_prompt_2, height=height, width=width, num_inference_steps=steps, guidance_scale=scale, num_images_per_prompt=1, generator=generator).images
    torch.cuda.empty_cache()
    refined_images = pipe2(prompt=prompt, prompt_2=prompt_2, image=int_images, guidance_scale=guidance_scale, aesthetic_score=aesthetic_score, negative_aesthetic_score=negative_aesthetic_score).images
    int_image_np = np.array(int_images[0])
    refined_image_np = np.array(refined_images[0])

    # Add watermark to the images
    int_image_np_with_watermark = add_watermark(int_image_np)
    refined_image_np_with_watermark = add_watermark(refined_image_np)

    # Save the generated images to Hugging Face Spaces
    save_image_to_hf_space(int_image_np_with_watermark, "int_image")
    save_image_to_hf_space(refined_image_np_with_watermark, "refined_image")

    return int_image_np_with_watermark, refined_image_np_with_watermark

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. While FFusion AI and Source Code Bulgaria Ltd. maintain the rights to their modifications and enhancements, all rights to the original Software and associated intellectual property remain with Stability AI Ltd. Use of the FFusionXL-BASE model is subject to the terms of this License, as well as any terms and conditions set forth by Stability AI Ltd. for the use of the original Software.</p>
    <p>Attribution: SDXL 0.9 is licensed under the SDXL Research License, Copyright (c) Stability AI Ltd. All Rights Reserved.</p>
        <h2>Warning and Compliance</h2>
    <p>Any use of the demo for generating inappropriate or unlawful content is strictly prohibited, and any misuse will not be tolerated. Individuals found to be generating content that is in violation of these terms or any applicable laws will be dealt with accordingly in accordance with legal regulations. The responsibility for any misuse or inappropriate use of the demo lies solely with the users who generated such content, and this demo, nor its affiliates, shall be held liable for any such use.</p>
    <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> </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='🎨 Main Prompt', placeholder='Describe the Desired Image (77 Token Limit)', lines=2), 
             gr.Textbox(label='❌ Negative Prompt', placeholder='Specify Unwanted Elements', lines=2),
             gr.Textbox(label='📝 Secondary Prompt for Text Encoder 2', placeholder='The prompt for tokenizer_2 and text_encoder_2 (Optional)', lines=1),
             gr.Textbox(label='❌ Negative Prompt for Text Encoder 2', placeholder='Negative guidance for text_encoder_2 (Optional)', lines=1),
             gr.Slider(3, 25, 10, label='🌊 DiFFusion Scale: Influence of Main Features'),
             gr.Slider(1, 10, 5, label='🧭 Guidance Scale: Intensity of Guidance'),
            #gr.Slider(640, 1024, 1024, step=128, label='📏 Height'),
            #gr.Slider(640, 1280, 1024, step=128, label='📐 Width'),
             gr.Slider(1, 10, 6, label='🎨 Aesthetic Score: Preference for Visual Appeal'),
             gr.Slider(1, 10, 2.5, label='🚫 Negative Aesthetic Score: Avoidance of Unwanted Aesthetics'),
             gr.Slider(10, maximum=80, value=50, step=1, label='💎 Number of Diffusion Steps'),
             gr.Slider(minimum=1, step=1, maximum=999999999999999999, randomize=True, label='🎲 Seed')],
   
             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;
                 }
             """,
             allow_flagging='never'
).launch(debug=True, max_threads=10)