app.py

import gradio as gr
import numpy as np
import random
import torch
import spaces
import os
import base64
import math

from PIL import Image
from diffusers import QwenImageEditPlusPipeline
from pillow_heif import register_heif_opener

from huggingface_hub import login
from prompt_augment import PromptAugment
login(token=os.environ.get('hf'))

register_heif_opener()

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = QwenImageEditPlusPipeline.from_pretrained(
    "FireRedTeam/FireRed-Image-Edit-1.1", 
    torch_dtype=dtype
).to(device)

prompt_handler = PromptAugment()

ADAPTER_SPECS = {
    "Covercraft": {
        "repo": "FireRedTeam/FireRed-Image-Edit-LoRA-Zoo",
        "weights": "FireRed-Image-Edit-Covercraft.safetensors",
        "adapter_name": "covercraft",
    },
    "Lightning": {
        "repo": "FireRedTeam/FireRed-Image-Edit-LoRA-Zoo",
        "weights": "FireRed-Image-Edit-Lightning-8steps-v1.0.safetensors",
        "adapter_name": "lightning",
    },
    "Makeup": {
        "repo": "FireRedTeam/FireRed-Image-Edit-LoRA-Zoo",
        "weights": "FireRed-Image-Edit-Makeup.safetensors",
        "adapter_name": "makeup",
    }
}

LOADED_ADAPTERS = set()
LORA_OPTIONS = ["None"] + list(ADAPTER_SPECS.keys())


def load_lora(lora_name):
    """加载并激活指定的 LoRA"""
    if lora_name == "None" or not lora_name:
        if LOADED_ADAPTERS:
            pipe.set_adapters([], adapter_weights=[])
        return
    
    spec = ADAPTER_SPECS.get(lora_name)
    if not spec:
        raise gr.Error(f"LoRA 配置未找到: {lora_name}")

    adapter_name = spec["adapter_name"]

    if adapter_name not in LOADED_ADAPTERS:
        print(f"--- Downloading and Loading Adapter: {lora_name} ---")
        try:
            pipe.load_lora_weights(
                spec["repo"], 
                weight_name=spec["weights"], 
                adapter_name=adapter_name
            )
            LOADED_ADAPTERS.add(adapter_name)
        except Exception as e:
            raise gr.Error(f"Failed to load adapter {lora_name}: {e}")
    else:
        print(f"--- Adapter {lora_name} is already loaded ---")

    pipe.set_adapters([adapter_name], adapter_weights=[1.0])


MAX_SEED = np.iinfo(np.int32).max
MAX_INPUT_IMAGES = 3


def limit_images(images):
    if images is None:
        return None
    if len(images) > MAX_INPUT_IMAGES:
        gr.Info(f"最多支持 {MAX_INPUT_IMAGES} 张图片，已自动移除多余图片")
        return images[:MAX_INPUT_IMAGES]
    return images


def calculate_dimensions(target_area, ratio):
    width = math.sqrt(target_area * ratio)
    height = width / ratio
    width = round(width / 32) * 32
    height = round(height / 32) * 32
    return int(width), int(height)


def update_dimensions_on_upload(images, max_area=1024*1024):
    if images is None or len(images) == 0:
        return 0, 0
    
    try:
        first_item = images[0]
        if isinstance(first_item, tuple):
            img = first_item[0]
        else:
            img = first_item
        
        if isinstance(img, Image.Image):
            pil_img = img
        elif isinstance(img, str):
            pil_img = Image.open(img)
        else:
            return 0, 0
        
        h, w = pil_img.height, pil_img.width
        is_multi_image = len(images) > 1
        
        if not is_multi_image:
            return 0, 0
        
        ratio = w / h
        new_w, new_h = calculate_dimensions(max_area, ratio)
        return new_h, new_w
    except Exception as e:
        print(f"获取图片尺寸失败: {e}")
        return 0, 0


@spaces.GPU(duration=180)
def infer(
    input_images,
    prompt,
    lora_choice,
    seed=42,
    randomize_seed=False,
    true_guidance_scale=4.0,
    num_inference_steps=40,
    height=None,
    width=None,
    rewrite_prompt=False,
    num_images_per_prompt=1,
    progress=gr.Progress(track_tqdm=True),
):
    negative_prompt = " "
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)
    
    load_lora(lora_choice)
    pil_images = []
    if input_images is not None:
        for item in input_images[:MAX_INPUT_IMAGES]:
            try:
                if isinstance(item, tuple):
                    img = item[0]
                else:
                    img = item
                    
                if isinstance(img, Image.Image):
                    pil_images.append(img.convert("RGB"))
                elif isinstance(img, str):
                    pil_images.append(Image.open(img).convert("RGB"))
            except Exception as e:
                print(f"处理图片出错: {e}")
                continue
    
    if height == 0:
        height = None
    if width == 0:
        width = None
    
    if rewrite_prompt and len(pil_images) > 0:
        prompt = prompt_handler.predict(prompt, [pil_images[0]])
        print(f"Rewritten Prompt: {prompt}")
    
    if pil_images:
        for i, img in enumerate(pil_images):
            print(f"    [{i}] size: {img.width}x{img.height}")
    images = pipe(
        image=pil_images if len(pil_images) > 0 else None,
        prompt=prompt,
        height=height,
        width=width,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=1.0,
        true_cfg_scale=true_guidance_scale,
        num_images_per_prompt=num_images_per_prompt,
    ).images

    return images, seed

css = """
#col-container { margin: 0 auto; max-width: 1200px; }
#edit-btn { height: 100% !important; min-height: 42px; }
"""



def get_image_base64(image_path):
    with open(image_path, "rb") as img_file:
        return base64.b64encode(img_file.read()).decode('utf-8')


logo_base64 = get_image_base64("logo.png") if os.path.exists("logo.png") else None

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        if logo_base64:
            gr.HTML(f'<img src="data:image/png;base64,{logo_base64}" alt="FireRed Logo" width="400" style="display: block; margin: 0 auto;">')
        else:
            gr.Markdown("# FireRed Image Edit")
        gr.Markdown(f"[Learn more](https://github.com/FireRedTeam/FireRed-Image-Edit) about the FireRed-Image-Edit series. Supports multi-image input (up to {MAX_INPUT_IMAGES} images.)")
        with gr.Row():
            with gr.Column(scale=1):
                input_images = gr.Gallery(
                    label="Upload Images",
                    type="pil",
                    interactive=True,
                    height=300,
                    columns=3,
                    object_fit="contain",
                )
            
            with gr.Column(scale=1):
                result = gr.Gallery(
                    label="Output Images",
                    type="pil",
                    height=300,
                    columns=2,
                    object_fit="contain",
                )
        
        prompt = gr.Textbox(
            label="Edit Prompt",
            placeholder="e.g., transform into anime..",
        )
        
        with gr.Row(equal_height=True):
            with gr.Column(scale=5):
                lora_choice = gr.Dropdown(
                    label="Choose Lora",
                    choices=LORA_OPTIONS,
                    value=LORA_OPTIONS[0] if LORA_OPTIONS else "None",
                )
            with gr.Column(scale=4):
                run_button = gr.Button("Edit Image", variant="primary", elem_id="edit-btn")

        with gr.Accordion("Advanced Settings", open=False):
            with gr.Row():
                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
            
            with gr.Row():
                true_guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=4.0)
                num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=40)
            
            with gr.Row():
                height = gr.Slider(label="Height (0=auto)", minimum=0, maximum=2048, step=8, value=0)
                width = gr.Slider(label="Width (0=auto)", minimum=0, maximum=2048, step=8, value=0)
            
            with gr.Row():
                rewrite_prompt = gr.Checkbox(label="Rewrite Prompt", value=False)
                num_images_per_prompt = gr.Slider(label="Num Images", minimum=1, maximum=4, step=1, value=1)

        # Examples
        gr.Examples(
            examples=[
                [["examples/master1.png"], "将背景换为带自然光效的浅蓝色，身穿浅米色蕾丝领上衣，将发型改为右侧佩戴精致珍珠发夹，同时单手向前抬起握着一把宝剑，另一只手自然摆放。面部微笑。",  "None"],
                [["examples/master2.png"], "替换背景为盛开的樱花树场景；更换衣服为黑色西装，为人物添加单肩蓝色书包，单手抓住包带。头发变为高马尾。色调明亮。蹲下。",  "None"],
                [["examples/master3_1.png", "examples/master3_2.png"], "把图1中的模特换成图2里的长裙和高帮帆布鞋，保持原有姿态和配饰，整体风格统一。",  "None"],
                [["examples/master4_1.png", "examples/master4_2.png"], "把图1中的白色衬衫和棕色半裙，换成图2里的灰褐色连帽卫衣、黑色侧边条纹裤、卡其色工装靴和同色云朵包，保持模特姿态和背景不变。",  "None"],
                [["examples/makeup1.png"], "为人物添加纯欲厌世妆：使用冷白皮哑光粉底均匀肤色，描绘细挑的灰黑色野生眉，眼部晕染浅灰调眼影并加深眼尾，画出上扬的黑色眼线，粘贴浓密卷翘的假睫毛，在眼头和卧蚕处提亮，涂抹深紫调哑光口红并勾勒唇形，在颧骨处扫上浅粉腮红，鼻梁和眉骨处打高光，下颌线处轻扫阴影。", "Makeup"],
                [["examples/makeup2.png"], "为人物添加妆容：使用象牙白哑光粉底均匀肤色，描绘细长柳叶眉并填充浅棕色，眼部晕染浅棕色眼影并加深眼尾，画出自然黑色眼线，粘贴浓密假睫毛，用浅棕色眼影提亮卧蚕；涂抹豆沙色哑光口红并勾勒唇形，在两颊扫上浅粉色腮红，在鼻梁和颧骨处轻扫高光，在面部轮廓处轻扫阴影。", "Makeup"],
                [["examples/text1_1.png", "examples/text1_2.png"], "请在图1添加主标题文本 “谁说我们丑了”，字体样式参考图2中主标题《人！给我开个罐罐》；主标题整体采用横向排版多行错落（非严格对齐），置于图片左下角；在狗狗右下方、贴近前爪附近添加一个手绘“爱心”涂鸦贴纸；增加鱼眼镜头效果", "Covercraft"],
                [["examples/text2_1.png", "examples/text2_2.png"], "请在图1添加主标题文本 “崽子第一次玩冰”，副标题“坐标：东南休闲公园”，主标题和副标题的字体样式参考图2中主标题“无露营不冬天”，主标题整体采用横向排版多行，主标题添加在画面左侧上方；副标题添加在画面左侧下方，字的层级更小，避免修改和遮挡图1主体关键信息（人物/核心景物）和画面中心。", "Covercraft"],
            ],
            inputs=[input_images, prompt, lora_choice],
            outputs=[result, seed],
            fn=infer,
            cache_examples=False,
            label="Examples"
        )

    # 监听 LoRA 选择变化：Lightning 时锁定参数
    def on_lora_change(lora_name):
        if lora_name == "Lightning":
            return (
                gr.update(value=8, interactive=False),      # num_inference_steps
                gr.update(value=1.0, interactive=False),    # true_guidance_scale
                gr.update(value=0, interactive=True),      # seed
                gr.update(value=False, interactive=False),  # randomize_seed
            )
        else:
            return (
                gr.update(value=40, interactive=True),      # num_inference_steps
                gr.update(value=4.0, interactive=True),     # true_guidance_scale
                gr.update(value=42, interactive=True),      # seed
                gr.update(value=True, interactive=True),    # randomize_seed
            )
    
    lora_choice.change(
        fn=on_lora_change,
        inputs=[lora_choice],
        outputs=[num_inference_steps, true_guidance_scale, seed, randomize_seed],
    )
    def on_image_upload(images):
        limited = limit_images(images)
        h, w = update_dimensions_on_upload(limited)
        return limited, h, w
    
    input_images.upload(
        fn=on_image_upload,
        inputs=[input_images],
        outputs=[input_images, height, width],
    )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            input_images,
            prompt, lora_choice, seed, randomize_seed,
            true_guidance_scale, num_inference_steps,
            height, width, rewrite_prompt, num_images_per_prompt,
        ],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    demo.queue()
    demo.launch(allowed_paths=["./"])