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='',
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) \
\
\
\
\
 \
\
').launch(debug=True, max_threads=10)
\
\
\
\
\
\