app.py

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
    )