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 re
import json
# 登录 Hugging Face Hub
HF_TOKEN = os.environ.get("HF_TOKEN")
login(token=HF_TOKEN)
import diffusers
print(diffusers.__version__)

# 初始化
dtype = torch.float16  # 您可以根据需要调整数据类型
device = "cuda" if torch.cuda.is_available() else "cpu"
base_model = "black-forest-labs/FLUX.1-dev"  # 替换为您的模型

# 加载管道
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()
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        self.end_time = time.time()
        self.elapsed_time = self.end_time - self.start_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")

# 定义位置、偏移和区域的映射
valid_locations = {  # x, y in 90*90
    'in the center': (45, 45),
    'on the left': (15, 45),
    'on the right': (75, 45),
    'on the top': (45, 15),
    'on the bottom': (45, 75),
    'on the top-left': (15, 15),
    'on the top-right': (75, 15),
    'on the bottom-left': (15, 75),
    'on the bottom-right': (75, 75)
}

valid_offsets = {  # x, y in 90*90
    'no offset': (0, 0),
    'slightly to the left': (-10, 0),
    'slightly to the right': (10, 0),
    'slightly to the upper': (0, -10),
    'slightly to the lower': (0, 10),
    'slightly to the upper-left': (-10, -10),
    'slightly to the upper-right': (10, -10),
    'slightly to the lower-left': (-10, 10),
    'slightly to the lower-right': (10, 10)
}

valid_areas = {  # w, h in 90*90
    "a small square area": (50, 50),
    "a small vertical area": (40, 60),
    "a small horizontal area": (60, 40),
    "a medium-sized square area": (60, 60),
    "a medium-sized vertical area": (50, 80),
    "a medium-sized horizontal area": (80, 50),
    "a large square area": (70, 70),
    "a large vertical area": (60, 90),
    "a large horizontal area": (90, 60)
}

# 解析角色位置的函数
def parse_character_position(character_position):
    # 定义正则表达式模式
    location_pattern = '|'.join(re.escape(key) for key in valid_locations.keys())
    offset_pattern = '|'.join(re.escape(key) for key in valid_offsets.keys())
    area_pattern = '|'.join(re.escape(key) for key in valid_areas.keys())

    # 提取位置
    location_match = re.search(location_pattern, character_position, re.IGNORECASE)
    location = location_match.group(0) if location_match else 'in the center'

    # 提取偏移
    offset_match = re.search(offset_pattern, character_position, re.IGNORECASE)
    offset = offset_match.group(0) if offset_match else 'no offset'

    # 提取区域
    area_match = re.search(area_pattern, character_position, re.IGNORECASE)
    area = area_match.group(0) if area_match else 'a medium-sized square area'

    return {
        'location': location,
        'offset': offset,
        'area': area
}

# 创建掩码的函数
def create_attention_mask(image_width, image_height, location, offset, area):
    # 图像在生成时通常会被缩放为 90x90，因此先定义一个基础尺寸
    base_size = 90

    # 获取位置坐标
    loc_x, loc_y = valid_locations.get(location, (45, 45))
    # 获取偏移量
    offset_x, offset_y = valid_offsets.get(offset, (0, 0))
    # 获取区域大小
    area_width, area_height = valid_areas.get(area, (60, 60))

    # 计算最终位置
    final_x = loc_x + offset_x
    final_y = loc_y + offset_y

    # 将坐标和尺寸映射到实际图像尺寸
    scale_x = image_width / base_size
    scale_y = image_height / base_size

    center_x = final_x * scale_x
    center_y = final_y * scale_y
    width = area_width * scale_x
    height = area_height * scale_y

    # 计算左上角和右下角坐标
    x_start = int(max(center_x - width / 2, 0))
    y_start = int(max(center_y - height / 2, 0))
    x_end = int(min(center_x + width / 2, image_width))
    y_end = int(min(center_y + height / 2, image_height))

    # 创建掩码
    mask = torch.zeros((image_height, image_width), dtype=torch.float32, device="cuda")
    mask[y_start:y_end, x_start:x_end] = 1.0

    # 展平成一维
    mask_flat = mask.view(-1)  # 形状为 (image_height * image_width,)
    return mask_flat

# 自定义注意力处理器

class CustomCrossAttentionProcessor(AttnProcessor2_0):
    def __init__(self, masks, adapter_names):
        super().__init__()
        self.masks = masks  # 列表，包含每个角色的掩码 (shape: [key_length])
        self.adapter_names = adapter_names  # 列表，包含每个角色的 LoRA 适配器名称

    def __call__(
        self,
        attn,
        hidden_states,
        encoder_hidden_states=None,
        attention_mask=None,
        temb=None,
        **kwargs,
    ):
        """
        自定义的注意力处理器，用于在注意力计算中应用角色掩码。

        参数：
            attn: 注意力模块实例。
            hidden_states: 输入的隐藏状态 (query)。
            encoder_hidden_states: 编码器的隐藏状态 (key/value)。
            attention_mask: 注意力掩码。
            temb: 时间嵌入（可能不需要）。
            **kwargs: 其他参数。

        返回：
            处理后的隐藏状态。
        """
        # 获取当前的 adapter_name
        adapter_name = getattr(attn, 'adapter_name', None)
        if adapter_name is None or adapter_name not in self.adapter_names:
            # 如果没有 adapter_name，或者不在我们的列表中，直接执行父类的 __call__ 方法
            return super().__call__(attn, hidden_states, encoder_hidden_states, attention_mask, temb, **kwargs)
        
        # 查找 adapter_name 对应的索引
        idx = self.adapter_names.index(adapter_name)
        mask = self.masks[idx]  # 获取对应的掩码 (shape: [key_length])

        # 以下是 AttnProcessor2_0 的实现，我们在适当的位置加入自定义的掩码逻辑

        residual = hidden_states
        if attn.spatial_norm is not None:
            hidden_states = attn.spatial_norm(hidden_states, temb)

        input_ndim = hidden_states.ndim

        if input_ndim == 4:
            batch_size, channel, height, width = hidden_states.shape
            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
        else:
            batch_size, sequence_length, _ = hidden_states.shape

        if encoder_hidden_states is None:
            encoder_hidden_states = hidden_states
        else:
            # 如果有 encoder_hidden_states，获取其形状
            encoder_batch_size, key_length, _ = encoder_hidden_states.shape

        if attention_mask is not None:
            # 处理 attention_mask，如果需要的话
            attention_mask = attn.prepare_attention_mask(attention_mask, key_length, batch_size)
            # attention_mask 的形状应为 (batch_size, attn.heads, query_length, key_length)
            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
        else:
            # 如果没有 attention_mask，我们创建一个全 0 的掩码
            attention_mask = torch.zeros(
                batch_size, attn.heads, 1, key_length, device=hidden_states.device, dtype=hidden_states.dtype
            )

        if attn.group_norm is not None:
            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)

        query = attn.to_q(hidden_states)

        if attn.norm_cross:
            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)

        key = attn.to_k(encoder_hidden_states)
        value = attn.to_v(encoder_hidden_states)

        inner_dim = key.shape[-1]
        head_dim = inner_dim // attn.heads

        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)

        if attn.norm_q is not None:
            query = attn.norm_q(query)
        if attn.norm_k is not None:
            key = attn.norm_k(key)

        # 计算原始的注意力得分
        # 我们需要在计算注意力得分前应用掩码
        # 但由于 PyTorch 的 scaled_dot_product_attention 接受 attention_mask 参数，我们需要调整我们的掩码

        # 创建自定义的 attention_mask
        # mask 的形状为 [key_length]，需要调整为 (batch_size, attn.heads, 1, key_length)
        custom_attention_mask = mask.view(1, 1, 1, -1).to(hidden_states.device, dtype=hidden_states.dtype)
        # 将有效位置设为 0，被掩蔽的位置设为 -1e9（对于 float16，使用 -65504）
        mask_value = -65504.0 if hidden_states.dtype == torch.float16 else -1e9
        custom_attention_mask = (1.0 - custom_attention_mask) * mask_value  # 有效位置为 0，无效位置为 -1e9

        # 将自定义掩码添加到 attention_mask
        attention_mask = attention_mask + custom_attention_mask

        # 计算注意力
        hidden_states = F.scaled_dot_product_attention(
            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
        )

        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
        hidden_states = hidden_states.to(query.dtype)

        # linear proj
        hidden_states = attn.to_out[0](hidden_states)
        # dropout
        hidden_states = attn.to_out[1](hidden_states)

        if input_ndim == 4:
            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)

        if attn.residual_connection:
            hidden_states = hidden_states + residual

        hidden_states = hidden_states / attn.rescale_output_factor

        return hidden_states


# 替换注意力处理器的函数
def replace_attention_processors(pipe, masks, adapter_names):
    custom_processor = CustomCrossAttentionProcessor(masks, adapter_names)
    for name, module in pipe.transformer.named_modules():
        if hasattr(module, 'attn'):
            module.attn.adapter_name = getattr(module, 'adapter_name', None)
            module.attn.processor = custom_processor
        if hasattr(module, 'cross_attn'):
            module.cross_attn.adapter_name = getattr(module, 'adapter_name', None)
            module.cross_attn.processor = custom_processor

# 生成图像的函数

def generate_image_with_embeddings(prompt_embeds, pooled_prompt_embeds, steps, seed, cfg_scale, width, height, progress):
    pipe.to(device)
    generator = torch.Generator(device=device).manual_seed(seed)

    with calculateDuration("Generating image"):
        # Generate image
        generated_image = pipe(
            prompt_embeds=prompt_embeds,
            pooled_prompt_embeds=pooled_prompt_embeds,
            num_inference_steps=steps,
            guidance_scale=cfg_scale,
            width=width,
            height=height,
            generator=generator,
        ).images[0]
    
    progress(99, "Generate success!")
    return generated_image

# 主函数
@spaces.GPU
@torch.inference_mode()
def run_lora(prompt_bg, character_prompts_json, character_positions_json, lora_strings_json, prompt_details, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
    
    # 解析角色提示词、位置和 LoRA 字符串
    try:
        character_prompts = json.loads(character_prompts_json)
        character_positions = json.loads(character_positions_json)
        lora_strings = json.loads(lora_strings_json)
    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON input: {e}")
    
    # 确保提示词、位置和 LoRA 字符串的数量一致
    if len(character_prompts) != len(character_positions) or len(character_prompts) != len(lora_strings):
        raise ValueError("The number of character prompts, positions, and LoRA strings must be the same.")
    
    # 角色的数量
    num_characters = len(character_prompts)
    
    # Load LoRA weights
    with calculateDuration("Loading LoRA weights"):
        pipe.unload_lora_weights()
        adapter_names = []
        for lora_info in lora_strings:
            lora_repo = lora_info.get("repo")
            weights = lora_info.get("weights")
            adapter_name = lora_info.get("adapter_name")
            if lora_repo and weights and adapter_name:
                # 调用 pipe.load_lora_weights() 方法加载权重
                pipe.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
                adapter_names.append(adapter_name)
                # 将 adapter_name 设置为模型的属性
                setattr(pipe.transformer, 'adapter_name', adapter_name)

            else:
                raise ValueError("Invalid LoRA string format. Each item must have 'repo', 'weights', and 'adapter_name' keys.")
        adapter_weights = [lora_scale] * len(adapter_names)
        # 调用 pipeline.set_adapters 方法设置 adapter 和对应权重
        pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
    
    # 确保 adapter_names 的数量与角色数量匹配
    if len(adapter_names) != num_characters:
        raise ValueError("The number of LoRA adapters must match the number of characters.")
    
    # Set random seed for reproducibility
    if randomize_seed:
        with calculateDuration("Set random seed"):
            seed = random.randint(0, MAX_SEED)
    
    # 编码提示词
    with calculateDuration("Encoding prompts"):
        # 编码背景提示词
        bg_text_input = pipe.tokenizer(prompt_bg, return_tensors="pt").to(device)
        bg_prompt_embeds = pipe.text_encoder_2(bg_text_input.input_ids.to(device))[0]
        bg_pooled_embeds = pipe.text_encoder(bg_text_input.input_ids.to(device)).pooler_output

        # 编码角色提示词
        character_prompt_embeds = []
        character_pooled_embeds = []
        for prompt in character_prompts:
            char_text_input = pipe.tokenizer(prompt, return_tensors="pt").to(device)
            char_prompt_embeds = pipe.text_encoder_2(char_text_input.input_ids.to(device))[0]
            char_pooled_embeds = pipe.text_encoder(char_text_input.input_ids.to(device)).pooler_output
            character_prompt_embeds.append(char_prompt_embeds)
            character_pooled_embeds.append(char_pooled_embeds)

        # 编码互动细节提示词
        details_text_input = pipe.tokenizer(prompt_details, return_tensors="pt").to(device)
        details_prompt_embeds = pipe.text_encoder_2(details_text_input.input_ids.to(device))[0]
        details_pooled_embeds = pipe.text_encoder(details_text_input.input_ids.to(device)).pooler_output

        # 合并背景和互动细节的嵌入
        prompt_embeds = torch.cat([bg_prompt_embeds, details_prompt_embeds], dim=1)
        pooled_prompt_embeds = torch.cat([bg_pooled_embeds, details_pooled_embeds], dim=1)
    
    # 解析角色位置
    character_infos = []
    for position_str in character_positions:
        info = parse_character_position(position_str)
        character_infos.append(info)
    
    # 创建角色的掩码
    masks = []
    for info in character_infos:
        mask = create_attention_mask(width, height, info['location'], info['offset'], info['area'])
        masks.append(mask)
    
    # 替换注意力处理器
    replace_attention_processors(pipe, masks, adapter_names)
    
    # Generate image
    final_image = generate_image_with_embeddings(prompt_embeddings, pooled_prompt_embeds, steps, seed, cfg_scale, width, height, progress)
    
    # 您可以在此处添加上传图片的代码
    result = {"status": "success", "message": "Image generated"}

    progress(100, "Completed!")

    return final_image, seed, json.dumps(result)

# Gradio 界面
css="""
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown("Flux with LoRA")
    with gr.Row():
        
        with gr.Column():

            prompt_bg = gr.Text(label="Background Prompt", placeholder="Enter background/scene prompt", lines=2)
            character_prompts = gr.Text(label="Character Prompts (JSON List)", placeholder='["Character 1 prompt", "Character 2 prompt"]', lines=5)
            character_positions = gr.Text(label="Character Positions (JSON List)", placeholder='["Character 1 position", "Character 2 position"]', lines=5)
            lora_strings_json = gr.Text(label="LoRA Strings (JSON List)", placeholder='[{"repo": "lora_repo1", "weights": "weights1", "adapter_name": "adapter_name1"}, {"repo": "lora_repo2", "weights": "weights2", "adapter_name": "adapter_name2"}]', lines=5)
            prompt_details = gr.Text(label="Interaction Details", placeholder="Enter interaction details between characters", lines=2)
            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)
                    lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=1, step=0.01, value=0.5)

                with gr.Row():
                    width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=512)
                    height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=512)

                with gr.Row():
                    cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=7.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_bg,
        character_prompts,
        character_positions,
        lora_strings_json,
        prompt_details,
        cfg_scale,
        steps,
        randomize_seed,
        seed,
        width,
        height,
        lora_scale,
        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()