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'FireRed Logo') 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=["./"])