app.py

import os
import gradio as gr
import numpy as np
import random
import spaces
import torch
import json
import logging
from diffusers import DiffusionPipeline
from huggingface_hub import login
import time
from datetime import datetime
from io import BytesIO
# from diffusers.models.attention_processor import AttentionProcessor
from diffusers.models.attention_processor import AttnProcessor2_0
import torch.nn.functional as F
import time
import boto3
from io import BytesIO
import re
import json

# Login Hugging Face Hub
HF_TOKEN = os.environ.get("HF_TOKEN")
login(token=HF_TOKEN)
import diffusers
print(diffusers.__version__)

# init
dtype = torch.float16  # use float16 for fast generate
device = "cuda" if torch.cuda.is_available() else "cpu"
base_model = "black-forest-labs/FLUX.1-dev"

# load pipe
pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)

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)
@torch.inference_mode()
def generate_image(prompt, steps, seed, cfg_scale, width, height, progress):
    
    with calculateDuration("Generating image"):
        # Generate image
        pipe.to(device)
        generator = torch.Generator(device=device).manual_seed(seed)
        print(prompt, steps, seed, cfg_scale, width, height)
        generated_image = pipe(
            prompt=prompt,
            num_inference_steps=steps,
            guidance_scale=cfg_scale,
            width=width,
            height=height,
            generator=generator,
        ).images[0]
    
    progress(99, "Generate success!")
    return generated_image


def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
    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)
    
    return image_file

def run_lora(prompt, lora_strings_json,  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)
    
    # Load LoRA weights
    if lora_strings_json:
        try:
            lora_configs = json.loads(lora_strings_json)
        except:
            lora_configs = None
    if lora_configs:
        with calculateDuration("Loading LoRA weights"):
            pipe.unload_lora_weights()
            adapter_names = []
            adapter_weights = []
            for lora_info in lora_configs:
                    lora_repo = lora_info.get("repo")
                    weights = lora_info.get("weights")
                    adapter_name = lora_info.get("adapter_name")
                    adapter_weight = lora_info.get("adapter_weight")
                    if lora_repo and weights and adapter_name:
                        # load lora
                        pipe.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
                        adapter_names.append(adapter_name)
                        adapter_weights.append(adapter_weight)
            # set lora weights
            pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
    
    # Set random seed for reproducibility
    if randomize_seed:
        with calculateDuration("Set random seed"):
            seed = random.randint(0, MAX_SEED)
    
    # Generate image
    error_message = ""
    try:
        final_image = generate_image(prompt, steps, seed, cfg_scale, width, height, progress)
    except Exception as e:
        error_message =  str(e)
        final_image = None
        
    if final_image:
        if upload_to_r2:
            with calculateDuration("Upload image"):
                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}
        
    progress(100, "Completed!")

    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)
            
            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():
                    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")

    inputs = [
        prompt,
        lora_strings_json,
        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
    )

demo.queue().launch()