app.py

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, variant="fp16", 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)

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

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

    return int_image_np, refined_image_np

   
gr.Interface(fn=genie, 
             inputs=[gr.Textbox(label='Describe your FFusion idea. 77 Token Limit.'), 
                     gr.Textbox(label='Things the AI should not create (negative prompt)'), 
                     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="Generated Image"), gc.Image(type='numpy', label="Refined Image")],
             title="FFusionXL - Generate and Refine", 
             description='<div style="display: flex; flex-wrap: wrap; gap: 2px;"><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 = '**Citation** \
                        Please note that the demo is intended for academic and research purposes ONLY. Any use of the demo for generating inappropriate content is strictly prohibited. The responsibility for any misuse or inappropriate use of the demo lies solely with the users who generated such content, and this demo shall not be held liable for any such use. Original code: Manjushri. By interacting within this environment, you hereby acknowledge and agree to the terms of the SDXL 0.9 Research License. \
                        Attribution: SDXL 0.9 is licensed under the SDXL Research License, Copyright (c) Stability AI Ltd. All Rights Reserved. \
                        **License** \
                        [SDXL 0.9 Research License](https://huggingface.co/stabilityai/stable-diffusion-xl-base-0.9/blob/main/LICENSE.md) \
                        [FFXL 0.9 Research License](https://huggingface.co/FFusion/FFusionXL-09-SDXL/blob/main/LICENSE.md) \
                        <div style="display: flex; flex-wrap: wrap; gap: 2px;">\
                        <img src="https://img.shields.io/badge/%F0%9F%94%A5%20Refiner%20Compatible-Yes-success"> \
                        <img src="https://img.shields.io/badge/%F0%9F%92%BB%20CLIP--ViT%2FG%20and%20CLIP--ViT%2FL%20tested-Yes-success"> \
                        <img src="https://img.shields.io/badge/%F0%9F%A7%A8%20FFXL%20Diffusers-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>\
                         \
                        ![ffusionAI-FFusionXL-SDXL-preview.jpg](https://cdn-uploads.huggingface.co/production/uploads/6380cf05f496d57325c12194/LIONhgnyxUuEynNivGsME.jpeg) \
                         \
                        <a href="mailto:[email protected]"><img src="https://img.shields.io/badge/Email-di%40ffusion.ai-blue?style=for-the-badge&logo=gmail"></a>').launch(debug=True, max_threads=10)