中山大学、AIGCの大規模応用を促進するためにソース拡散モデル統合コードフレームワークを公開

近年、拡散モデルに基づく画像生成モデルが次々と登場し、驚くべき生成効果を示しています。しかし、関連する研究モデルの既存のコードフレームワークは過度に断片化されており、統一されたフレームワークシステムが欠如しているため、「移行が難しい」、「敷居が高い」、「品質が低い」などのコード実装の困難が生じています。

この目的のために、中山大学のヒューマン・コンピュータ・オブジェクト融合研究所（HCP ラボ）は、モデルの微調整、パーソナライズされたトレーニング、推論の最適化、画像編集など、拡散モデルに基づく関連アルゴリズムを体系的に実装した HCP-Diffusion フレームワークを構築しました。その構造を図 1 に示します。

図 1 HCP-Diffusion フレームワーク構造図。これは、統一されたフレームワークを通じて既存の拡散関連手法を統合し、さまざまなモジュール式のトレーニングおよび推論最適化手法を提供します。

HCP-Diffusion は、統一された形式の構成ファイルを通じてさまざまなコンポーネントとアルゴリズムを展開し、フレームワークの柔軟性とスケーラビリティを大幅に向上させます。開発者は、コード実装の詳細を繰り返すことなく、ビルディング ブロックのようにアルゴリズムを組み合わせることができます。

たとえば、HCP-Diffusion をベースに、構成ファイルを変更するだけで、LoRA、DreamArtist、ControlNet などのさまざまな一般的なアルゴリズムの展開と組み合わせを完了できます。これにより、イノベーションのハードルが下がるだけでなく、フレームワークがさまざまなカスタマイズされたデザインと互換性を持つようになります。

• HCP-Diffusion コードツール: https://github.com/7eu7d7/HCP-Diffusion
• HCP-Diffusion グラフィカル インターフェース: https://github.com/7eu7d7/HCP-Diffusion-webui

HCP-拡散: 機能モジュールの紹介

フレームワークの機能

HCP-Diffusion は、現在主流の拡散トレーニング アルゴリズム フレームワークをモジュール化して、フレームワークの普遍性を実現します。主な機能は次のとおりです。

• 統一アーキテクチャ: Diffusion シリーズのモデル用の統一コード フレームワークの構築
• オペレータ プラグイン: データ、トレーニング、推論、パフォーマンス最適化 (deepspeed、colossal-AI、オフロード アクセラレーション最適化など) などのオペレータ アルゴリズムをサポートします。
• ワンクリック構成: Diffusion シリーズ モデルは、構成ファイルを変更することで、高い柔軟性をもって実装できます。
• ワンクリックトレーニング: Web UI、ワンクリックトレーニング、推論機能を提供します

データモジュール

HCP-Diffusion は、それぞれ異なる画像サイズと注釈形式を使用できる複数の並列データセットの定義をサポートしています。図 2 に示すように、各トレーニング反復では、トレーニング用に各データセットからバッチを抽出します。さらに、各データ セットは複数のデータ ソースで構成でき、txt、json、yaml などの注釈形式やカスタム注釈形式をサポートし、非常に柔軟なデータ前処理および読み込みメカニズムを備えています。

図2 データセット構造の概略図

データセット処理部分は、自動クラスタリングを備えたアスペクト比バケットを提供し、異なる画像サイズのデータ​​セットの処理をサポートします。ユーザーはデータセットのサイズに対して追加の処理や調整を行う必要はありません。フレームワークは、アスペクト比や解像度に基づいて最適なグループ化方法を自動的に選択します。このテクノロジーにより、データ処理の閾値が大幅に下がり、ユーザーエクスペリエンスが最適化され、開発者はアルゴリズム自体の革新にさらに集中できるようになります。

画像データの前処理に関しては、トーチビジョンやアルブメンテーションなどのさまざまな画像処理ライブラリとも互換性があります。ユーザーは必要に応じて設定ファイルで前処理方法を直接設定したり、これに基づいてカスタマイズされた画像処理方法を拡張したりできます。

図3.データセット構成ファイルの例

テキスト注釈に関しては、HCP-Diffusion は、複雑で多様なトレーニング方法とデータ注釈をサポートできる、柔軟で明確なプロンプト テンプレート仕様を設計しました。上記の設定ファイルのソース ディレクトリにある word_names を適用します。ここでは、下の図の中括弧内の特殊文字に対応する埋め込み単語ベクトルとカテゴリの説明をカスタマイズして、DreamBooth や DreamArtist などのモデルと互換性を持たせることができます。

図4 プロンプトテンプレート

さらに、テキスト注釈については、TagDropout や TagShuffle などのさまざまなテキスト強化方法が提供されており、画像とテキスト データ間の過剰適合の問題を軽減し、生成される画像の多様性を高めます。

モデルフレームワークモジュール

HCP-Diffusion は、現在主流の拡散トレーニング アルゴリズム フレームワークをモジュール化することで、フレームワークの普遍性を実現します。具体的には、イメージ エンコーダーとイメージ デコーダーがイメージのエンコードとデコードを完了し、ノイズ ジェネレーターがフォワード プロセスのノイズを生成し、拡散モデルが拡散プロセスを実装し、条件エンコーダーが生成条件をエンコードし、アダプタがモデルを微調整して下流のタスクと一致させ、正と負のデュアル チャネルが正と負の条件によるイメージの制御生成を表します。

図 5. モデル構造の構成例 (モデル プラグイン、カスタム ワードなど)

図 5 に示すように、HCP-Diffusion は、構成ファイル内の単純な組み合わせを通じて、LoRA、ControlNet、DreamArtist などの複数の主流のトレーニング アルゴリズムを実装できます。また、LoRA と Textual Inversion を同時にトレーニングしたり、独自のトリガーワードを LoRA にバインドしたりするなど、上記のアルゴリズムの組み合わせもサポートしています。さらに、プラグイン モジュールを使用すると、現在の主流のアクセス方法すべてと互換性のあるプラグインを簡単にカスタマイズできます。上記のモジュール化により、HCP-Diffusion はあらゆる主流アルゴリズムのフレームワーク構築を実現し、開発の敷居を下げ、モデルの共同イノベーションを促進します。

HCP-Diffusion は、LoRA や ControlNet などのさまざまなアダプタ クラスのアルゴリズムをモデル プラグインに抽象化します。共通のモデル プラグイン ベース クラスをいくつか定義することで、このようなアルゴリズムをすべて均一に扱うことができ、ユーザーの使用と開発コストを削減し、すべてのアダプタ クラスのアルゴリズムを統一することができます。

このフレームワークは、現在の主流アルゴリズムをすべて簡単にサポートできる 4 種類のプラグインを提供します。

+ SinglePluginBlock: lora シリーズなど、そのレイヤーの入力に基づいて出力を変更するシングルレイヤー プラグイン。挿入レイヤーを定義するための正規表現 (re: プレフィックス) をサポートしていますが、pre_hook: プレフィックスはサポートしていません。

+ PluginBlock: 残差接続を定義する場合など、入力層と出力層は 1 つだけです。挿入レイヤーを定義するための正規表現 (re: プレフィックス) をサポートし、入力レイヤーと出力レイヤーの両方が pre_hook: プレフィックスをサポートします。

+ MultiPluginBlock: 入力層と出力層の両方に、controlnet などの複数のブロックを設定できます。正規表現 (re: プレフィックス) はサポートされていません。入力層と出力層の両方で pre_hook: プレフィックスがサポートされています。

+ WrapPluginBlock: 元のモデルのレイヤーを置き換え、元のモデルのレイヤーをこのクラスのオブジェクトにします。置換レイヤーを定義するための正規表現 (re: プレフィックス) はサポートされていますが、pre_hook: プレフィックスはサポートされていません。

トレーニングと推論モジュール

図6 カスタムオプティマイザ構成

HCP-Diffusion の設定ファイルは、実行時に自動的にインスタンス化される Python オブジェクトの定義をサポートしています。この設計により、開発者は、上の図に示すように、フレームワーク コードを変更することなく、カスタム オプティマイザー、損失関数、ノイズ サンプラーなどの pip でインストール可能なカスタム モジュールに簡単にアクセスできます。設定ファイルは構造が明確で理解しやすく、再現性も高いため、学術研究とエンジニアリング展開をスムーズにつなげることができます。

加速最適化サポート

HCP-Diffusion は、Accelerate、DeepSpeed、Colossal-AI などの複数のトレーニング最適化フレームワークをサポートしており、トレーニング中のメモリ使用量を大幅に削減し、トレーニングを高速化できます。 EMA 操作をサポートすることで、モデルの生成効果と一般化がさらに向上します。推論フェーズでは、モデルのオフロードや VAE タイリングなどの操作がサポートされ、最小 1 GB のビデオ メモリで画像生成を完了できます。

図7 モジュラー構成ファイル

上記のシンプルなファイル構成により、上図に示すように、関連するフレームワーク リソースを見つけるために多大な労力を費やすことなく、モデルの構成を完了できます。 HCP-Diffusion のモジュール設計により、モデル メソッドの定義、トレーニング ロジック、推論ロジックが完全に分離されます。モデルを構成するときに、トレーニング部分と推論部分のロジックを考慮する必要がないため、ユーザーはメソッド自体に集中できます。同時に、HCP-Diffusion は、ほとんどの主流アルゴリズムのフレームワーク構成例を提供しており、いくつかのパラメータを変更するだけで展開できます。

HCP-Diffusion: Web UI グラフィカル インターフェース

HCP-Diffusion は、構成ファイルを直接変更できることに加えて、画像生成やモデルトレーニングなどの複数のモジュールを含む対応する Web UI グラフィックインターフェイスを提供し、ユーザーエクスペリエンスを向上させ、フレームワークの学習しきい値を大幅に削減し、アルゴリズムの理論から実践への変換を加速します。

図8 HCP-Diffusion Web UIイメージインターフェース

研究室紹介

中山大学のヒューマン・コンピュータ・オブジェクト融合研究所（HCPラボ）は、2010年に林静教授によって設立されました。近年、マルチモーダルコンテンツ理解、因果および認知推論、具体化学習などの分野で豊富な学術成果を達成し、国内外の科学技術賞や最優秀論文賞を数多く受賞しており、製品レベルのAI技術とプラットフォームの構築に取り組んでいます。研究室ウェブサイト: http://www.sysu-hcp.net