import os
import gradio as gr
import numpy as np
import random
import spaces
import torch
import json
import logging
from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL, AutoPipelineForImage2Image
from huggingface_hub import login
from diffusers.utils import load_image
from lora_loading_patch import load_lora_into_transformer
import time
from datetime import datetime
from io import BytesIO
import torch.nn.functional as F
from PIL import Image, ImageFilter
import time
import boto3
from io import BytesIO
import re
import json
import random
import string
# Login Hugging Face Hub
HF_TOKEN = os.environ.get("HF_TOKEN")
login(token=HF_TOKEN)
import diffusers
# init
dtype = torch.bfloat16
device = "cuda"
base_model = "black-forest-labs/FLUX.1-dev"
# load pipe
taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
txt2img_pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
# txt2img_pipe.__class__.load_lora_into_transformer = classmethod(load_lora_into_transformer)
# img2img model
img2img_pipe = AutoPipelineForImage2Image.from_pretrained(base_model, vae=good_vae, transformer=txt2img_pipe.transformer, text_encoder=txt2img_pipe.text_encoder, tokenizer=txt2img_pipe.tokenizer, text_encoder_2=txt2img_pipe.text_encoder_2, tokenizer_2=txt2img_pipe.tokenizer_2, torch_dtype=dtype)
# img2img_pipe.__class__.load_lora_into_transformer = classmethod(load_lora_into_transformer)
MAX_SEED = 2**32 - 1
class calculateDuration:
def __init__(self, activity_name=""):
self.activity_name = activity_name
def __enter__(self):
self.start_time = time.time()
self.start_time_formatted = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(self.start_time))
print(f"Activity: {self.activity_name}, Start time: {self.start_time_formatted}")
return self
def __exit__(self, exc_type, exc_value, traceback):
self.end_time = time.time()
self.elapsed_time = self.end_time - self.start_time
self.end_time_formatted = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(self.end_time))
if self.activity_name:
print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
else:
print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
@spaces.GPU(duration=120)
def generate_image(orginal_image, prompt, adapter_names, steps, seed, image_strength, cfg_scale, width, height, progress):
gr.Info("Start to generate images ...")
with calculateDuration(f"Make a new generator: {seed}"):
txt2img_pipe.to(device)
generator = torch.Generator(device=device).manual_seed(seed)
with calculateDuration("Generating image"):
# Generate image
joint_attention_kwargs = {"scale": 1}
if orginal_image:
generated_image = img2img_pipe(
prompt=prompt,
image=orginal_image,
strength=image_strength,
num_inference_steps=steps,
guidance_scale=cfg_scale,
width=width,
height=height,
generator=generator,
joint_attention_kwargs=joint_attention_kwargs
).images[0]
else:
generated_image = txt2img_pipe(
prompt=prompt,
num_inference_steps=steps,
guidance_scale=cfg_scale,
width=width,
height=height,
max_sequence_length=512,
generator=generator,
joint_attention_kwargs=joint_attention_kwargs
).images[0]
torch.cuda.empty_cache()
progress(99, "Generate image success!")
return generated_image
def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
with calculateDuration("Upload images"):
print("upload_image_to_r2", account_id, access_key, secret_key, bucket_name)
connectionUrl = f"https://{account_id}.r2.cloudflarestorage.com"
s3 = boto3.client(
's3',
endpoint_url=connectionUrl,
region_name='auto',
aws_access_key_id=access_key,
aws_secret_access_key=secret_key
)
current_time = datetime.now().strftime("%Y/%m/%d/%H%M%S")
image_file = f"generated_images/{current_time}_{random.randint(0, MAX_SEED)}.png"
buffer = BytesIO()
image.save(buffer, "PNG")
buffer.seek(0)
s3.upload_fileobj(buffer, bucket_name, image_file)
print("upload finish", image_file)
# start to generate thumbnail
thumbnail = image.copy()
thumbnail_width = 256
aspect_ratio = image.height / image.width
thumbnail_height = int(thumbnail_width * aspect_ratio)
thumbnail = thumbnail.resize((thumbnail_width, thumbnail_height), Image.LANCZOS)
# Generate the thumbnail image filename
thumbnail_file = image_file.replace(".png", "_thumbnail.png")
# Save thumbnail to buffer and upload
thumbnail_buffer = BytesIO()
thumbnail.save(thumbnail_buffer, "PNG")
thumbnail_buffer.seek(0)
s3.upload_fileobj(thumbnail_buffer, bucket_name, thumbnail_file)
print("upload thumbnail finish", thumbnail_file)
return image_file
def generate_random_4_digit_string():
return ''.join(random.choices(string.digits, k=4))
def run_lora(prompt, image_url, lora_strings_json, image_strength, cfg_scale, steps, randomize_seed, seed, width, height, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
print("run_lora", prompt, lora_strings_json, cfg_scale, steps, width, height)
gr.Info("Starting process")
torch.cuda.empty_cache()
img2img_model = False
orginal_image = None
if image_url:
orginal_image = load_image(image_url)
img2img_model = True
# Set random seed for reproducibility
if randomize_seed:
with calculateDuration("Set random seed"):
seed = random.randint(0, MAX_SEED)
# Load LoRA weights
gr.Info("Start to load LoRA ...")
with calculateDuration("Unloading LoRA"):
img2img_pipe.unload_lora_weights()
# img2img_pipe.unload_lora()
txt2img_pipe.unload_lora_weights()
txt2img_pipe.unload_lora()
print(txt2img_pipe.get_active_adapters())
list_adapters_component_wise = txt2img_pipe.get_list_adapters()
print(list_adapters_component_wise)
lora_configs = None
adapter_names = []
lora_names = []
if lora_strings_json:
try:
lora_configs = json.loads(lora_strings_json)
except:
gr.Warning("Parse lora config json failed")
print("parse lora config json failed")
if lora_configs:
with calculateDuration("Loading LoRA weights"):
adapter_weights = []
for idx, lora_info in enumerate(lora_configs):
lora_repo = lora_info.get("repo")
weights = lora_info.get("weights")
adapter_name = lora_info.get("adapter_name")
lora_name = generate_random_4_digit_string()
lora_names.append(lora_name)
adapter_weight = lora_info.get("adapter_weight")
adapter_names.append(adapter_name)
adapter_weights.append(adapter_weight)
if lora_repo and weights and adapter_name:
try:
if img2img_model:
img2img_pipe.load_lora_weights(lora_repo, weight_name=weights, low_cpu_mem_usage=True, adapter_name=lora_name)
img2img_pipe.set
else:
txt2img_pipe.load_lora_weights(lora_repo, weight_name=weights, low_cpu_mem_usage=True, adapter_name=lora_name)
except:
print("load lora error")
# set lora weights
if len(lora_names) > 0:
if img2img_model:
img2img_pipe.set_adapters(lora_names, adapter_weights=adapter_weights)
else:
txt2img_pipe.set_adapters(lora_names, adapter_weights=adapter_weights)
print(txt2img_pipe.get_active_adapters())
# Generate image
error_message = ""
try:
print("Start applying for zeroGPU resources")
final_image = generate_image(orginal_image, prompt, adapter_names, steps, seed, image_strength, cfg_scale, width, height, progress)
except Exception as e:
error_message = str(e)
gr.Error(error_message)
print("Run error", e)
final_image = None
if final_image:
if upload_to_r2:
url = upload_image_to_r2(final_image, account_id, access_key, secret_key, bucket)
result = {"status": "success", "message": "upload image success", "url": url}
else:
result = {"status": "success", "message": "Image generated but not uploaded"}
else:
result = {"status": "failed", "message": error_message}
gr.Info("Completed！")
progress(100, "Completed!")
torch.cuda.empty_cache()
return final_image, seed, json.dumps(result)
# Gradio interface
css="""
#col-container {
margin: 0 auto;
max-width: 640px;
}
"""
with gr.Blocks(css=css) as demo:
gr.Markdown("flux-dev-multi-lora")
with gr.Row():
with gr.Column():
prompt = gr.Text(label="Prompt", placeholder="Enter prompt", lines=10)
lora_strings_json = gr.Text(label="LoRA Configs (JSON List String)", placeholder='[{"repo": "lora_repo1", "weights": "weights1", "adapter_name": "adapter_name1", "adapter_weight": 1}, {"repo": "lora_repo2", "weights": "weights2", "adapter_name": "adapter_name2", "adapter_weight": 1}]', lines=5)
image_url = gr.Text(label="Image url", placeholder="Enter image url to enable image to image model", lines=1)
run_button = gr.Button("Run", scale=0)
with gr.Accordion("Advanced Settings", open=False):
with gr.Row():
seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
with gr.Row():
width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
with gr.Row():
image_strength = gr.Slider(label="Denoise Strength", info="Lower means more image influence", minimum=0.1, maximum=1.0, step=0.01, value=0.75)
cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
upload_to_r2 = gr.Checkbox(label="Upload to R2", value=False)
account_id = gr.Textbox(label="Account Id", placeholder="Enter R2 account id")
access_key = gr.Textbox(label="Access Key", placeholder="Enter R2 access key here")
secret_key = gr.Textbox(label="Secret Key", placeholder="Enter R2 secret key here")
bucket = gr.Textbox(label="Bucket Name", placeholder="Enter R2 bucket name here")
with gr.Column():
result = gr.Image(label="Result", show_label=False)
seed_output = gr.Text(label="Seed")
json_text = gr.Text(label="Result JSON")
gr.Markdown("**Disclaimer:**")
gr.Markdown(
"This demo is only for research purpose. This space owner cannot be held responsible for the generation of NSFW (Not Safe For Work) content through the use of this demo. Users are solely responsible for any content they create, and it is their obligation to ensure that it adheres to appropriate and ethical standards. This space owner provides the tools, but the responsibility for their use lies with the individual user."
)
inputs = [
prompt,
image_url,
lora_strings_json,
image_strength,
cfg_scale,
steps,
randomize_seed,
seed,
width,
height,
upload_to_r2,
account_id,
access_key,
secret_key,
bucket
]
outputs = [result, seed_output, json_text]
run_button.click(
fn=run_lora,
inputs=inputs,
outputs=outputs
)
try:
demo.queue().launch()
except:
print("demo exception ...")