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, 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
from diffusers import FluxPipeline
from huggingface_hub import hf_hub_download
from diffusers.quantizers import PipelineQuantizationConfig
from diffusers import (FluxPriorReduxPipeline, FluxInpaintPipeline, FluxFillPipeline, FluxKontextPipeline, FluxPipeline)

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

# init
dtype = torch.bfloat16
device = "cuda:0"

base_model = "black-forest-labs/FLUX.1-Krea-dev"

# pipeline_quant_config = PipelineQuantizationConfig(
#     quant_backend="bitsandbytes_4bit",
#     quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
#     components_to_quantize=["transformer", "text_encoder_2"],
# )

txt2img_pipe = FluxKontextPipeline.from_pretrained(base_model, torch_dtype=dtype)
txt2img_pipe = txt2img_pipe.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")
        
def safe_trim_for_clip(text: str, max_words: int = 77) -> str:
    # 简单按词裁，不破坏主 prompt。你也可以做更智能的关键词抽取。
    tokens = re.split(r"\s+", text.strip())
    if len(tokens) <= max_words:
        return text
    return " ".join(tokens[:max_words])

def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
    with calculateDuration("Upload images"):
        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))

@spaces.GPU(duration=120)
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")
    pipe = txt2img_pipe
    device = pipe.device
    print(device)

    # ========== Seed ==========
    if randomize_seed:
        with calculateDuration("Set random seed"):
            seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(seed)

    # ========== LoRA ==========
    gr.Info("Start to load LoRA ...")
    with calculateDuration("Unloading LoRA"):
        try:
            pipe.unload_lora_weights()
        except Exception as _:
            # 某些版本上未加载时调用可能抛异常，忽略
            pass

    adapter_names = []
    adapter_weights = []
    if lora_strings_json:
        try:
            lora_configs = json.loads(lora_strings_json)
        except Exception as _:
            lora_configs = None
            gr.Warning("Parse lora config json failed")
            print("parse lora config json failed")

        if lora_configs:
            with calculateDuration("Loading LoRA weights"):
                for lora_info in lora_configs:
                    repo = lora_info.get("repo")
                    weights = lora_info.get("weights")
                    # 优先使用用户提供的 adapter_name；没有则随机
                    adapter_name = lora_info.get("adapter_name") or f"adp_{generate_random_4_digit_string()}"
                    weight = float(lora_info.get("adapter_weight", 1.0))
                    if not (repo and weights):
                        print(f"skip invalid lora entry: {lora_info}")
                        continue
                    try:
                        weight_path = hf_hub_download(repo_id=repo, filename=weights)
                        # 关键修复：prefix=None，避免仅在 text_encoder 查找
                        pipe.load_lora_weights(weight_path, adapter_name=adapter_name, prefix=None)
                        adapter_names.append(adapter_name)
                        adapter_weights.append(weight)
                    except Exception as e:
                        print(f"load lora error for {repo}/{weights}: {e}")

                if adapter_names:
                    pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
                    # 可选：融合后推理更快，但无法动态调整权重
                    # pipe.fuse_lora(adapter_names=adapter_names)
                    try:
                        active = pipe.get_active_adapters() if hasattr(pipe, "get_active_adapters") else []
                        print("Active adapters:", active)
                    except Exception as e:
                        print("Active adapters query failed:", e)
                    
                    lora_layer_count = 0
                    for name, module in pipe.transformer.named_modules():
                        attrs = dir(module)
                        if any(a.startswith("lora_") for a in attrs) or "lora" in module.__class__.__name__.lower():
                            lora_layer_count += 1
                    print(f"[DEBUG] transformer LoRA layers: {lora_layer_count}")
                    
                    # 若层数为 0，给出直观警告
                    if lora_layer_count == 0:
                        gr.Warning("LoRA seems not injected (0 layers on transformer). Check whether the LoRA is trained for FLUX and `prefix=None` is set.")


    pipe.enable_vae_slicing()
    clip_side_prompt = safe_trim_for_clip(prompt, max_words=77) 
    init_image = None
    error_message = ""
    try:
        gr.Info("Start to generate images ...")
        joint_attention_kwargs = {"scale": 1}
        image = pipe(
            prompt=prompt,
            num_inference_steps=int(steps),
            guidance_scale=float(cfg_scale),
            width=int(width),
            height=int(height),
            max_sequence_length=512,
            generator=generator,
            joint_attention_kwargs=joint_attention_kwargs
        ).images[0]
    except Exception as e:
        error_message = str(e)
        gr.Error(error_message)
        print("fatal error", e)
        image = None

    result = {"status": "failed", "message": error_message} if image is None else {"status": "success", "message": "Image generated but not uploaded"}
    if image is not None and upload_to_r2:
        try:
            url = upload_image_to_r2(image, account_id, access_key, secret_key, bucket)
            result = {"status": "success", "message": "upload image success", "url": url}
        except Exception as e:
            err = f"Upload failed: {e}"
            gr.Warning(err)
            print(err)
            result = {"status": "success", "message": "generated but upload failed"}

    gr.Info("Completed！")
    progress(100, "Completed!")
    return json.dumps(result)

# Gradio interface



with gr.Blocks() 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():
            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 = [json_text]

    run_button.click(
        fn=run_lora,
        inputs=inputs,
        outputs=outputs
    )
    
try:
    demo.queue().launch()
except:
   print("demo exception ...")