英語版で読む
技術約7分で読めます

Qwen-Image 2.1のTaylorSeer修正がDiffusers本体にマージされた

Qwen画像生成DiffusersTaylorSeerOSSGitHubPyTorchPythonClaude Code
いけさん目次

Qwen-Image 2.1をTaylorSeerキャッシュで速くした記事では、既定のKVキャッシュと併用すると2ステップ目でエラーになり、スクリプトの中で TaylorSeerState.update を差し替えて回避した。
その回避策をDiffusersにissueとして報告し、同じ内容をPRにしたところ、9月25日にmainブランチへマージされた。
issueを立てたのが9月21日の夜なので、4日弱で取り込まれた。

元の不具合

Qwen-Image 2.1のパイプラインは、既定でKVキャッシュ（use_kv_cache=True）が有効になっている。
文章と参照画像の部分は1ステップ目で計算して保存し、2ステップ目以降は生成する画像の部分だけを計算する。
そのため、Transformerの出力の長さが1ステップ目だけ長い。512×512の生成なら、1ステップ目が文章側14トークンと画像1024トークンの計1038トークン、2ステップ目以降は1024トークンになる。

TaylorSeerは前後のステップの出力の差を取って次のステップの値を予測するので、長さの違う2つの差分を計算できずエラーになっていた。

RuntimeError: The size of tensor a (1024) must match the size of tensor b (1038) at non-singleton dimension 1

use_kv_cache=False にすれば動くが、参照画像を使う編集では参照画像の部分を毎ステップ計算し直すことになり、M1 Maxで832×1216・40ステップの編集が277.6秒から484.4秒に延びた。

issue作成からPR提出まで

9月21日22時17分にissue #14829を立てた。
再現コードやエラーログ、回避策のモンキーパッチ、M1 Maxでの計測値を載せ、確かめたのはApple Silicon（MPS）とTaylorSeerだけだと書いた。
モデル側で直す方がきれいかもしれない、という案も添えた。Qwen-Image 2.1のパイプラインは1ステップ目の出力のうち文章側を捨てているので、最後の出力層に画像のトークンだけを通せば長さはそろう。

38分後にメンテナのSayak Paul氏から、PRを出してもらえないかと返信があった。
その日の23時45分にPR #14831を出した。

修正とPRの文面はClaude Codeを使って作った。
Diffusersのコントリビューションガイドには、AIエージェントを使ったPRについての節がある。PRの説明に書く内容なども含めて要件が並んでいて、今回の進め方に関わった要件は表のとおり。

求められていること今回やったこと
PRの前にissueで方針を相談し、メンテナの明示的な了承を待つissueでPRを出してほしいと返信をもらってから出した
リポジトリにある self-review スキルで差分を確認するClaude Codeで self-review スキルを実行した
self-review の結果をPRの説明かコメントに載せるPRの説明に、未解決の点と判定を載せた

self-review はリポジトリの .ai/skills/ に置かれたエージェント用のスキルで、差分を、CIで動いているClaudeによる自動レビューと同じ基準に照らして確認する。
今回の結果は、修正が必要な問題はなし、判定はREADYだった。
レビューに回した未解決の点は、パイプライン全体を通すテストを入れていないこと、KVキャッシュ併用時に最初に実計算するステップ数（disable_cache_before_step）が2以上でないと予測を始められないこと、TaylorSeer以外のキャッシュは確かめていないこと。

取り込まれた修正

修正はDiffusersの src/diffusers/hooks/taylorseer_cache.py の1か所で、条件を1つ追加してコメントを3行付けただけ。

# The feature shape can change between steps, e.g. Qwen-Image 2.1 returns prefix + target
# tokens on the KV-cache prefill step and target tokens only afterwards. Stale factors
# cannot be differenced against the new features, so restart the expansion from order 0.
if prev is None or prev.shape != new_factors[j].shape:
    break
new_factors[j + 1] = (new_factors[j] - prev.to(features.dtype)) / delta_step

TaylorSeerは各ステップで、0次の値（出力そのもの）と、前のステップとの差から求める1次以上の値を保存している。
修正前は前の値が残っていれば必ず差を取っていた。修正後は形が違えば差を取らず、今回の出力を0次の値として保存するだけにする。形が変わらないときの動きは変わらない。

記事で使ったモンキーパッチは、形が変わったら保存していた値と最後に更新したステップをまとめて捨ててから元の処理を呼んでいた。
やり方は違うが、どちらも形が変わったステップを初回として扱う点は同じで、同じ条件で出した画像はファイルのMD5まで一致した。

レビューで求められたこと

修正の中身についての指摘はほとんどなく、動く証拠とテストを求められた。時刻は日本時間。

日時メンテナからの要望返したもの
9月22日18時台単体テストに「このテストが分からない」、修正の行に「これで期待どおり動くのか」。最小の再現コードと、期待される結果、短縮される時間も見せてほしい9月23日0時11分に、issueと同じ再現コード、mainでのエラー、計測の表
9月23日0時14分実際の出力画像も見せてほしい9月25日0時3分に、キャッシュなしとTaylorSeer＋KVキャッシュの画像を並べた表
9月25日12時42分FluxのTransformerにあるTaylorSeerのテストクラスと同じ形で、Qwen-Image 2.1用のテストを追加してほしい15時1分に、Qwen-Image 2.1のTransformerのテストファイルへテストクラスを追加

計測の表は、PRのブランチではなく同じ内容のモンキーパッチで測った値だと断ったうえで出した。出力が一致することは確かめてあったので、測り直してはいない。

最初に付けたテストは、TaylorSeerの内部状態に長さの違うテンソルを直接入力し、エラーにならず予測値が合うかを確かめるだけのテストだった。
ここで別のメンテナから、Fluxのテストのような形にできないかという提案があり、Sayak Paul氏が、Qwen-Image 2.1専用のテストクラスを追加するという意味だと補足した。

追加したテストクラスは、Fluxと同じくTaylorSeerの共通テスト一式を継承し、そこにKVキャッシュとの併用のテストを1つ追加した。
このテストは、小さな設定のQwen-Image 2.1のTransformerで、パイプラインと同じように1ステップ目で文章側を保存し、2ステップ目以降はKVキャッシュを使う8ステップのループを回す。
TaylorSeerを入れた状態で、KVキャッシュありとなしの出力が全ステップで一致するかを確かめる。修正前のmainブランチではissueと同じエラーが出て、修正後は通った。

TaylorSeerが予測に切り替わっていなければ、KVキャッシュありとなしが一致するのは当然で比較の意味がなくなるので、最後に4ステップ目の出力がキャッシュを使わない出力と違うことも確かめた。
この確認にはSayak Paul氏から「This is a very meaningful assertion!」とコメントが付いた。

15時31分に承認され、メンテナがPRのコメントからGPU上でこのテストを実行して通ったあと、16時4分にマージされた。

修正を使うには

9月26日の時点で、Diffusersの最新リリースは8月20日のv0.40.0で、この修正は入っていない。
次のリリースまでは、GitHubのmainブランチからDiffusersをインストールするか、前回の記事のモンキーパッチを使うことになる。

pip install git+https://github.com/huggingface/diffusers

修正はQwen-Image 2.1のモデル側ではなく、TaylorSeerのフック側に入った。なお、issueで触れたモデル側での直し方は、PRのレビューでは特に話題にならなかった。

関連記事

参考リンク