app.py

import gradio as gr
import numpy as np
import random, torch
import spaces  # [uncomment to use ZeroGPU]
from PIL import Image

from kontext.pipeline_flux_kontext import FluxKontextPipeline
from kontext.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
from diffusers import FluxTransformer2DModel
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

# ---------------------------
# utils
# ---------------------------
def resize_by_bucket(images_pil, resolution=512):
    assert len(images_pil) > 0, "images_pil 不能为空"
    bucket_override = [
        (336, 784), (344, 752), (360, 728), (376, 696),
        (400, 664), (416, 624), (440, 592), (472, 552),
        (512, 512),
        (552, 472), (592, 440), (624, 416), (664, 400),
        (696, 376), (728, 360), (752, 344), (784, 336),
    ]
    # 按目标分辨率缩放，并对齐到 16
    bucket_override = [(int(h / 512 * resolution), int(w / 512 * resolution)) for h, w in bucket_override]
    bucket_override = [(h // 16 * 16, w // 16 * 16) for h, w in bucket_override]

    aspect_ratios = [img.height / img.width for img in images_pil]
    mean_aspect_ratio = float(np.mean(aspect_ratios))

    new_h, new_w = bucket_override[0]
    min_aspect_diff = abs(new_h / new_w - mean_aspect_ratio)
    for h, w in bucket_override:
        aspect_diff = abs(h / w - mean_aspect_ratio)
        if aspect_diff < min_aspect_diff:
            min_aspect_diff = aspect_diff
            new_h, new_w = h, w

    resized_images = [img.resize((new_w, new_h), resample=Image.BICUBIC) for img in images_pil]
    return resized_images

# ---------------------------
# pipeline init
# ---------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"

flux_pipeline = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev")
flux_pipeline.scheduler = FlowMatchEulerDiscreteScheduler.from_config(flux_pipeline.scheduler.config)
flux_pipeline.scheduler.config.stochastic_sampling = False

# precision & device
flux_pipeline.vae.to(device).to(torch.bfloat16)
flux_pipeline.text_encoder.to(device).to(torch.bfloat16)
flux_pipeline.text_encoder_2.to(device).to(torch.bfloat16)

# 替换 transformer 权重
ckpt_path = hf_hub_download("NoobDoge/Multi_Ref_Model", "full_ema_model.safetensors")
flux_pipeline.transformer = FluxTransformer2DModel.from_single_file(ckpt_path, torch_dtype=torch.bfloat16)
flux_pipeline.transformer.to(device).to(torch.bfloat16)

# ---------------------------
# constants
# ---------------------------
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 512   # 与下方滑块默认值 1024 保持一致

# ---------------------------
# inference
# ---------------------------
@spaces.GPU  # [uncomment to use ZeroGPU]
def infer(
    prompt,
    ref1,            # PIL.Image 或 None
    ref2,            # PIL.Image 或 None（可选）
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,      # 目前没传入 pipeline，如需要可在下面调用里加上
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    # 组装可选参考图列表
    refs = [x for x in (ref1, ref2) if x is not None]
    if len(refs) == 0:
        raise gr.Error("请至少上传一张参考图（ref1 或 ref2）。")
    

    # 规范宽高：不超过 MAX_IMAGE_SIZE 且对齐到 16
    width = max(16, min(width, MAX_IMAGE_SIZE)) // 16 * 16
    height = max(16, min(height, MAX_IMAGE_SIZE)) // 16 * 16

    # 随机种子
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(int(seed))

    # 参考图按桶缩放
    raw_images = resize_by_bucket(refs, resolution=MAX_IMAGE_SIZE)

    # 推理
    with torch.no_grad():
        out = flux_pipeline(
            image=raw_images,
            prompt=prompt,
            height=height,
            width=width,
            num_inference_steps=int(num_inference_steps),
            max_area=MAX_IMAGE_SIZE ** 2,
            generator=generator,
            # 如需 guidance_scale，确保 pipeline 支持这个参数后再打开：
            # guidance_scale=float(guidance_scale),
        )
        output_img = out.images[0]

    return output_img, int(seed)

# ---------------------------
# UI
# ---------------------------
examples = [
    "Place the butterfly from the first image onto the landscape of the second image, positioning it either flying above the river near the bridge or perched on one of the trees in the foreground. Adjust the butterfly's size and blending to ensure it fits naturally in the scene.",
    "An astronaut riding a green horse",
    "A delicious ceviche cheesecake slice",
]

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

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# Multi Ref Edit Demo")

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button("Run", scale=0, variant="primary")

        # 两张输入图片（ref2 可空）
        with gr.Row():
            ref1_comp = gr.Image(label="Input Image 1", type="pil")
            ref2_comp = gr.Image(label="Input Image 2 (optional)", type="pil")

        result = gr.Image(label="Result", show_label=False)

        with gr.Accordion("Advanced Settings", open=False):
            seed_comp = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            randomize_seed_comp = gr.Checkbox(label="Randomize seed", value=True)

            with gr.Row():
                width_comp = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=512,
                )
                height_comp = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=512,
                )

            with gr.Row():
                guidance_scale_comp = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=10.0,
                    step=0.1,
                    value=2.5,
                )
                num_inference_steps_comp = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )

        gr.Examples(examples=[[e] for e in examples], inputs=[prompt])

    # 注意：不要把 [ref1, ref2] 当作列表传给 inputs！
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            ref1_comp,
            ref2_comp,              # ref2 可为空
            seed_comp,
            randomize_seed_comp,
            width_comp,
            height_comp,
            guidance_scale_comp,
            num_inference_steps_comp,
        ],
        outputs=[result, seed_comp],
    )

if __name__ == "__main__":
    demo.launch()