import os import gc import math import gradio as gr import numpy as np import spaces import torch import random from PIL import Image, ImageOps from typing import Iterable from gradio.themes import Soft from gradio.themes.utils import colors, fonts, sizes from datetime import datetime colors.orange_red = colors.Color( name="orange_red", c50="#FFF0E5", c100="#FFE0CC", c200="#FFC299", c300="#FFA366", c400="#FF8533", c500="#FF4500", c600="#E63E00", c700="#CC3700", c800="#B33000", c900="#992900", c950="#802200", ) class OrangeRedTheme(Soft): def __init__( self, *, primary_hue: colors.Color | str = colors.gray, secondary_hue: colors.Color | str = colors.orange_red, neutral_hue: colors.Color | str = colors.slate, text_size: sizes.Size | str = sizes.text_lg, font: fonts.Font | str | Iterable[fonts.Font | str] = ( fonts.GoogleFont("Outfit"), "Arial", "sans-serif", ), font_mono: fonts.Font | str | Iterable[fonts.Font | str] = ( fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace", ), ): super().__init__( primary_hue=primary_hue, secondary_hue=secondary_hue, neutral_hue=neutral_hue, text_size=text_size, font=font, font_mono=font_mono, ) super().set( background_fill_primary="*primary_50", background_fill_primary_dark="*primary_900", body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)", body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)", button_primary_text_color="white", button_primary_text_color_hover="white", button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)", button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)", button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)", button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)", button_secondary_text_color="black", button_secondary_text_color_hover="white", button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)", button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)", button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)", button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)", slider_color="*secondary_500", slider_color_dark="*secondary_600", block_title_text_weight="600", block_border_width="3px", block_shadow="*shadow_drop_lg", button_primary_shadow="*shadow_drop_lg", button_large_padding="11px", color_accent_soft="*primary_100", block_label_background_fill="*primary_200", ) orange_red_theme = OrangeRedTheme() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("CUDA_VISIBLE_DEVICES =", os.environ.get("CUDA_VISIBLE_DEVICES")) print("torch.__version__ =", torch.__version__) print("Using device:", device) from diffusers import FlowMatchEulerDiscreteScheduler from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3 dtype = torch.bfloat16 pipe = QwenImageEditPlusPipeline.from_pretrained( "FireRedTeam/FireRed-Image-Edit-1.0", transformer=QwenImageTransformer2DModel.from_pretrained( "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19", torch_dtype=dtype, device_map="cuda" ), torch_dtype=dtype ).to(device) try: pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3()) print("Flash Attention 3 Processor set successfully.") except Exception as e: print(f"Warning: Could not set FA3 processor: {e}") MAX_SEED = np.iinfo(np.int32).max def choose_safe_canvas_size(image, max_long_side=768, max_area=768 * 768, multiple=64): w, h = image.size area = w * h if area <= 0: return 512, 512 scale_by_side = max_long_side / max(w, h) scale_by_area = math.sqrt(max_area / area) scale = min(1.0, scale_by_side, scale_by_area) target_w = max(multiple, int(round((w * scale) / multiple)) * multiple) target_h = max(multiple, int(round((h * scale) / multiple)) * multiple) return target_w, target_h def prepare_input_image(image, padding_color=(255, 255, 255)): """ 【修复方案】：在 resize 前添加 protective padding (保护性边框) image: 输入的 PIL 图像 padding_color: 边框颜色，默认为白色 (RGB 255, 255, 255)。 你可以根据模型输出风格调整，比如黑色 (0,0,0) 或中灰 (128,128,128) """ # 1. 修复 EXIF 方向 (关键，防止手机照片颠倒) image = ImageOps.exif_transpose(image).convert("RGB") # 2. 计算并添加 protective padding # 我们根据原图尺寸，增加 10% 的安全边距 orig_w, orig_h = image.size pad_fraction = 0.20 # 增加 10% 的边框 pad_w = int(orig_w * pad_fraction) pad_h = int(orig_h * pad_fraction) # 创建新的大画布，填充指定颜色 new_canvas_size = (orig_w + 2 * pad_w, orig_h + 2 * pad_h) padded_image = Image.new("RGB", new_canvas_size, color=padding_color) # 将原图粘贴到画布中央 padded_image.paste(image, (pad_w, pad_h)) # 3. 计算用于网络输入的 Safe Canvas Size # 注意：使用 64 倍数对齐以获得最佳效果 target_w, target_h = choose_safe_canvas_size(padded_image, multiple=64) # 4. 将带有边框的图片缩放到目标尺寸 if padded_image.size != (target_w, target_h): final_input_image = padded_image.resize((target_w, target_h), Image.Resampling.LANCZOS) else: final_input_image = padded_image return final_input_image, target_w, target_h @spaces.GPU(duration=15) def run_gpu_inference(pil_image, prompt, seed, guidance_scale, steps, width, height): generator = torch.Generator(device=device).manual_seed(seed) negative_prompt = ( "worst quality, low quality, bad anatomy, bad hands, text, error, " "missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, " "signature, watermark, username, blurry" ) try: with torch.inference_mode(): kwargs = dict( image=[pil_image], prompt=prompt, height=height, width=width, num_inference_steps=steps, generator=generator, true_cfg_scale=guidance_scale, ) # true_cfg_scale <= 1 时不要传 negative_prompt，省掉无效分支和 warning if guidance_scale > 1: kwargs["negative_prompt"] = negative_prompt result_image = pipe(**kwargs).images[0] return result_image except Exception as e: raise e def process_and_concat(images, prompt, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress(track_tqdm=True)): if not images: raise gr.Error("Please upload at least one image to edit.") item = images[0] path_or_img = item[0] if isinstance(item, (tuple, list)) else item try: if isinstance(path_or_img, str): orig_image = Image.open(path_or_img).convert("RGB") elif isinstance(path_or_img, Image.Image): orig_image = path_or_img.convert("RGB") else: orig_image = Image.open(path_or_img.name).convert("RGB") orig_image = ImageOps.exif_transpose(orig_image).convert("RGB") except Exception as e: raise gr.Error(f"Could not load image: {e}") prepared_img, safe_w, safe_h = prepare_input_image(orig_image) if randomize_seed: seed = random.randint(0, MAX_SEED) raw_result_image = run_gpu_inference( prepared_img, prompt, seed, guidance_scale, steps, safe_w, safe_h ) gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # 不再把生成结果拉回“原图尺寸” # 只把原图缩到结果尺寸用于对比 final_result = raw_result_image orig_resized = orig_image.resize(final_result.size, Image.Resampling.LANCZOS) border_width = 24 frame_color = "#FFF0E5" total_width = final_result.width + (border_width * 2) total_height = final_result.height + orig_resized.height + (border_width * 3) concat_img = Image.new("RGB", (total_width, total_height), color=frame_color) concat_img.paste(final_result, (border_width, border_width)) concat_img.paste(orig_resized, (border_width, border_width * 2 + final_result.height)) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") filename = f"[edit]_{timestamp}_seed{seed}.png" filepath = os.path.join(os.getcwd(), filename) concat_img.save(filepath, format="PNG") return concat_img, filepath, seed css = """ #col-container { margin: 0 auto; max-width: 1000px; } #main-title h1 {font-size: 2.4em !important;} """ with gr.Blocks() as demo: with gr.Column(elem_id="col-container"): gr.Markdown("# FireRed-Image-Edit-1.0-Fast", elem_id="main-title") gr.Markdown("Perform image edits using [FireRed-Image-Edit-1.0](https://huggingface.co/FireRedTeam/FireRed-Image-Edit-1.0) with 4-step fast inference. Open on [GitHub](https://github.com/PRITHIVSAKTHIUR/FireRed-Image-Edit-1.0-Fast)") with gr.Row(equal_height=True): with gr.Column(): images = gr.Gallery( label="Upload Images", type="filepath", columns=2, rows=1, height=300, allow_preview=True ) prompt = gr.Text( label="Edit Prompt", show_label=True, max_lines=2, placeholder="e.g., transform into anime, upscale, change lighting...", ) run_button = gr.Button("Edit Image", variant="primary") with gr.Column(): output_image = gr.Image(interactive=False, format="png") output_file = gr.File(label="Download Lossless Merged Image") with gr.Accordion("Advanced Settings", open=False, visible=False): seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0) randomize_seed = gr.Checkbox(label="Randomize Seed", value=True) guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0) steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=4) gr.Markdown("[*](https://huggingface.co/FireRedTeam/FireRed-Image-Edit-1.0)This is still an experimental Space for FireRed-Image-Edit-1.0.") run_button.click( fn=process_and_concat, inputs=[images, prompt, seed, randomize_seed, guidance_scale, steps], outputs=[output_image, output_file, seed] ) if __name__ == "__main__": demo.queue(max_size=30).launch( css=css, theme=orange_red_theme, mcp_server=True, ssr_mode=False, show_error=True )