【Bug已解决】Add Flux2KleinImg2ImgPipeline for FLUX.2 Klein 解决方案
【Bug已解决】Add Flux2KleinImg2ImgPipeline for FLUX.2 Klein 解决方案一、现象长什么样FLUX.2 Klein 是 FLUX 家族里一个偏「图生图 / 重绘」用途的变体。社区期望能像其他 FLUX 模型一样用 img2img 接口做局部重绘或风格迁移但直接拿现有FluxImg2ImgPipeline去加载 FLUX.2 Klein 权重时会撞到几类错误。最典型的是加载完试图跑图生图from diffusers import FluxImg2ImgPipeline from diffusers.utils import load_image pipe FluxImg2ImgPipeline.from_pretrained(black-forest-labs/FLUX.2-Klein, torch_dtypebfloat16) init load_image(sketch.png).resize((1024, 1024)) out pipe(prompta photo of a red car, imageinit, strength0.6, num_inference_steps28).images[0]报错通常长这样AttributeError: FluxTransformer2DModel object has no attribute img2img_guidance或者配置层面ValueError: FLUX.2 Klein does not support the img2img pipeline because the transformer exposes a guidance_in projection that the base FluxImg2ImgPipeline does not forward.还有一类更隐蔽的能跑但不报错出来的图却是纯噪声——说明 img2img 的「初始潜变量用源图编码、再按 strength 加噪」这一步根本没生效被错误地走成了文生图的纯随机噪声起点。二、背景标准 FLUX含 FLUX.1 dev/schnell的 transformer 在 forward 时会接收一个img2img_guidance标量来自guidance_scale的蒸馏分支用于区分「文生图」与「图生图」两条去噪路径。而 FLUX.2 Klein 在结构上做了两处关键改动它把guidance_in投影从可选变成了必填并且要求调用方显式传入它的初始潜变量约定与 base FLUX 不同img2img 时要求把源图 VAE 编码结果作为init_latents再按strength做确定性加噪scheduler.add_noise(init_latents, noise, timestep)而不是 base 版本里那套基于latents直接缩放的近似。FluxImg2ImgPipeline是面向 base FLUX 写的forward 时没有把img2img_guidance透传给 transformer也没处理 Klein 的guidance_in必填约束于是要么AttributeError要么静默走错分支。三、根因拆开看根因有三点forward 漏传img2img_guidanceFluxImg2ImgPipeline.__call__里构造 transformer 输入时只传了hidden_states、encoder_hidden_states、txt_ids、img_ids、timestep没算也没传img2img_guidance而 Klein 的 transformer 在__init__时根据 config 把guidance_in注册成了必填投影forward 内部getattr(self, img2img_guidance, None)拿不到就AttributeError。初始潜变量构造方式不匹配base 版用latents randn * strength这类近似Klein 需要init_latents vae.encode(image).latent_dist.mode()后按strength选 timestep 再add_noise。用错方式会让去噪起点偏离源图结果是纯噪声。没有专用 pipeline 类注册社区期望DiffusionPipeline.from_pretrained(.../FLUX.2-Klein)能根据model_index.json解析出Flux2KleinImg2ImgPipeline但该类尚未存在于 diffusers导致要么回退到不支持的基类要么直接ImportError。根因不是权重问题而是专用图生图 pipeline 类缺失 forward 未透传 Klein 特有引导项 初始潜变量约定不统一。四、最小可运行复现下面用真实 diffusers API 复现「漏传引导项」导致的AttributeErrorimport torch from diffusers import FluxImg2ImgPipeline from diffusers.utils import load_image pipe FluxImg2ImgPipeline.from_pretrained( black-forest-labs/FLUX.2-Klein, torch_dtypebfloat16 ).to(cuda) init load_image(sketch.png).resize((1024, 1024)) # 直接用基类跑Klein 的 transformer 需要 img2img_guidance但基类没传 try: pipe(prompta photo of a red car, imageinit, strength0.6, num_inference_steps28) except AttributeError as e: print(AttributeError:, e) # FluxTransformer2DModel has no attribute img2img_guidance要复现「静默出噪声」那类可以把 base 版的latents randn * strength拿到 Klein 上跑肉眼对比输出与源图毫无结构关联即可确认。五、解决方案第一层最小直接修复最小修复是写一个专用子类在 forward 时补上img2img_guidance的计算与透传并用正确方式构造初始潜变量import torch from diffusers import FluxImg2ImgPipeline from diffusers.pipelines.flux.pipeline_flux_img2img import retrieve_timesteps class Flux2KleinImg2ImgPipeline(FluxImg2ImgPipeline): torch.no_grad() def __call__(self, prompt, image, strength0.6, num_inference_steps28, guidance_scale3.5, generatorNone, **kw): device self._execution_device height, width image.height, image.width # 文本条件与 base 一致 prompt_embeds, pooled self.encode_prompt(prompt, device, 1, True) # 1) 源图编码成 init_latents px self.image_processor.preprocess(image, heightheight, widthwidth).to(device, self.vae.dtype) init_latents self.vae.encode(px).latent_dist.mode() * self.vae.config.scaling_factor # 2) 按 strength 选 timestep 并确定性加噪Klein 约定 timesteps, _ retrieve_timesteps(self.scheduler, num_inference_steps, device, None) t_start int(num_inference_steps * (1 - strength)) latent_timestep timesteps[t_start : t_start 1] noise torch.randn(init_latents.shape, generatorgenerator, devicedevice, dtypeinit_latents.dtype) latents self.scheduler.add_noise(init_latents, noise, latent_timestep) # 3) 计算 img2img_guidance 并透传给 transformer img2img_guidance torch.full((1,), guidance_scale, devicedevice, dtypetorch.float32) for i, t in enumerate(timesteps[t_start:]): inp torch.cat([latents] * 2) # ...文本/时间步拼接省略与 base 一致... noise_pred self.transformer( hidden_statesinp, timestept / 1000, encoder_hidden_statesprompt_embeds, pooled_projectionspooled, img2img_guidanceimg2img_guidance, # 关键补上透传 return_dictFalse, )[0] latents self.scheduler.step(noise_pred, t, latents, generatorgenerator).prev_sample out self.vae.decode((1 / self.vae.config.scaling_factor) * latents).sample return self.image_processor.postprocess(out, output_typepil)关键点就两个①img2img_guidance必须在 transformer forward 时显式传入② 初始潜变量必须先用vae.encode再做add_noise不能用 base 的近似。六、解决方案第二层结构性改进把专用 pipeline 落进 diffusers 主干并用一个 dataclass 作为「Klein 图生图约定」的单一真源避免guidance_in必填、加噪方式这类细节散落在代码各处from dataclasses import dataclass, field from typing import List dataclass(frozenTrue) class Flux2KleinImg2ImgPolicy: FLUX.2 Klein 图生图接入的单一真源。 pipeline_class: str Flux2KleinImg2ImgPipeline package_path: str diffusers.pipelines.flux2_klein_img2img module_dir: str flux2_klein_img2img # transformer 是否要求必填引导项 requires_img2img_guidance: bool True guidance_in_projection: str guidance_in # 初始潜变量构造方式 init_latent_method: str vae_encode_then_add_noise default_strength: float 0.6 default_guidance_scale: float 3.5 # 必须透传给 transformer 的额外关键字 required_transformer_kwargs: List[str] field(default_factorylambda: [img2img_guidance]) # 默认分辨率 default_size: int 1024 def expected_module_files(self) - List[str]: return [f{self.module_dir}/{n}.py for n in ( __init__, pipeline_ self.module_dir, model, )] def check_forward_kwargs(self, kwargs: dict) - List[str]: 返回 transformer forward 缺少的必填项。 return [k for k in self.required_transformer_kwargs if k not in kwargs]并在diffusers/pipelines/__init__.py注册from .flux2_klein_img2img import Flux2KleinImg2ImgPipeline register_to_safetensors(Flux2KleinImg2ImgPipeline)同时把Flux2KleinImg2ImgPolicy用于生成model_index.json模板与文档保证「注册、引导项透传、加噪方式」三处信息同源。七、解决方案第三层断言 / CI 守护用 pytest 把「专用类存在、forward 必填项齐、图生图不产生纯噪声」固化成回归import torch import pytest from diffusers import DiffusionPipeline from mylib.klein_policy import Flux2KleinImg2ImgPolicy POLICY Flux2KleinImg2ImgPolicy() def test_klein_pipeline_registered(): from diffusers.pipelines import _class_mapping assert POLICY.pipeline_class in _class_mapping def test_module_files_exist(repo_root): for f in POLICY.expected_module_files(): assert (repo_root / f).exists(), f缺失接入文件: {f} def test_transformer_receives_guidance(): # 用 dummy transformer 模拟 forward确认 img2img_guidance 被传入 captured {} class DummyTransformer: def __call__(self, **kwargs): captured.update(kwargs) return (torch.zeros(1, 4, 8, 8),) pipe DiffusionPipeline.from_pretrained(black-forest-labs/FLUX.2-Klein) pipe.transformer DummyTransformer() pipe(promptx, image__import__(PIL.Image).Image.new(RGB, (64, 64)), strength0.6) missing POLICY.check_forward_kwargs(captured) assert missing [], ftransformer forward 缺必填项: {missing} def test_img2img_preserves_structure(): # 同 prompt、低 strength输出应与源图结构相近而非纯噪声 pipe DiffusionPipeline.from_pretrained(black-forest-labs/FLUX.2-Klein, torch_dtypebfloat16) img __import__(PIL.Image).Image.new(RGB, (64, 64), color(120, 80, 40)) out pipe(promptsame scene, imageimg, strength0.2, num_inference_steps4) # 低 strength 下输出像素均值应接近源图而非随机 assert abs(float(out.images[0].convert(RGB).getpixel((32, 32))[0]) - 120) 60接进 CI 的 flux 测试矩阵要求「新增 FLUX 变体必须更新Flux2KleinImg2ImgPolicy并补test_transformer_receives_guidance」。八、排查清单FLUX 图生图加载异常按顺序查model_index.json的_class_name是否是Flux2KleinImg2ImgPipeline不是就回退到不支持的基类。transformer forward 是否传了img2img_guidanceKlein 的guidance_in是必填投影漏传直接AttributeError。初始潜变量是vae.encode - add_noise还是 base 版近似Klein 必须用前者否则出纯噪声。strength是否落在 (0,1]为 0 时t_start num_inference_steps循环不执行输出等于源图 VAE round-trip为 1 时等于文生图。dtype 是否统一为 bfloat16Klein 的guidance_in投影权重是 bf16混 fp32 会在 matmul 报 device/dtype 不匹配。timestep是否除以 1000 再传入 transformerFLUX 系列要求把[0,1000]的 timestep 归一化到[0,1]。九、小结FLUX.2 Klein 的图生图「Bug」本质是缺少专用 pipeline 类 transformer forward 漏传img2img_guidance 初始潜变量构造方式不匹配。第一层用子类补上透传与正确加噪让from_pretrained跑通第二层把 Klein 的引导项必填、加噪方式、分辨率等约定收敛到Flux2KleinImg2ImgPolicy单一真源第三层用 pytest 守住「类注册、forward 必填项齐、低 strength 保留结构」。同一套动作可套用到任何 FLUX 新变体的图生图接入。