Sdxs

🚀 Simple Diffusion XS

XSサイズで、超品質

このモデルは、サイズが小さくても高品質な画像生成を実現します。消費者向けグラフィックカードでのトレーニングが可能で、多言語に対応しています。

🚀 クイックスタート

トレーニング状態（一時停止）: エポックN 16

💻 使用例

基本的な使用法

import torch
from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel
from transformers import AutoModel, AutoTokenizer
from PIL import Image
from tqdm.auto import tqdm
import os

def encode_prompt(prompt, negative_prompt, device, dtype):
    if negative_prompt is None:
        negative_prompt = ""

    with torch.no_grad():
        positive_inputs = tokenizer(
            prompt,
            return_tensors="pt",
            padding="max_length",
            max_length=512,
            truncation=True,
        ).to(device)
        positive_embeddings = text_model.encode_texts(
            positive_inputs.input_ids, positive_inputs.attention_mask
        )
        if positive_embeddings.ndim == 2:
            positive_embeddings = positive_embeddings.unsqueeze(1)
        positive_embeddings = positive_embeddings.to(device, dtype=dtype)
        
        negative_inputs = tokenizer(
            negative_prompt,
            return_tensors="pt",
            padding="max_length",
            max_length=512,
            truncation=True,
        ).to(device)
        negative_embeddings = text_model.encode_texts(negative_inputs.input_ids, negative_inputs.attention_mask)
        if negative_embeddings.ndim == 2:
            negative_embeddings = negative_embeddings.unsqueeze(1)
        negative_embeddings = negative_embeddings.to(device, dtype=dtype)
    return torch.cat([negative_embeddings, positive_embeddings], dim=0)

def generate_latents(embeddings, height=576, width=576, num_inference_steps=50, guidance_scale=5.5):
    with torch.no_grad():
        device, dtype = embeddings.device, embeddings.dtype
        half = embeddings.shape[0] // 2
        latent_shape = (half, 16, height // 8, width // 8)
        latents = torch.randn(latent_shape, device=device, dtype=dtype)
        embeddings = embeddings.repeat_interleave(half, dim=0)

        scheduler.set_timesteps(num_inference_steps)

        for t in tqdm(scheduler.timesteps, desc="Генерация"):
            latent_model_input = torch.cat([latents] * 2)
            latent_model_input = scheduler.scale_model_input(latent_model_input, t)
            noise_pred = unet(latent_model_input, t, embeddings).sample
            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
            latents = scheduler.step(noise_pred, t, latents).prev_sample
    return latents


def decode_latents(latents, vae, output_type="pil"):
    latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
    with torch.no_grad():
        images = vae.decode(latents).sample
    images = (images / 2 + 0.5).clamp(0, 1)
    images = images.cpu().permute(0, 2, 3, 1).float().numpy()
    if output_type == "pil":
        images = (images * 255).round().astype("uint8")
        images = [Image.fromarray(image) for image in images]
    return images

# Example usage:
if __name__ == "__main__":
    device = "cuda"
    dtype = torch.float16

    prompt = "кот"
    negative_prompt = "bad quality"
    tokenizer = AutoTokenizer.from_pretrained("visheratin/mexma-siglip")
    text_model = AutoModel.from_pretrained(
        "visheratin/mexma-siglip", torch_dtype=dtype, trust_remote_code=True
    ).to(device, dtype=dtype).eval()
    
    embeddings = encode_prompt(prompt, negative_prompt, device, dtype)    

    pipeid = "AiArtLab/sdxs"
    variant = "fp16"
    
    unet = UNet2DConditionModel.from_pretrained(pipeid, subfolder="unet", variant=variant).to(device, dtype=dtype).eval()
    vae = AutoencoderKL.from_pretrained(pipeid, subfolder="vae", variant=variant).to(device, dtype=dtype).eval()
    scheduler = DDPMScheduler.from_pretrained(pipeid, subfolder="scheduler")


    height, width = 576, 576
    num_inference_steps = 40
    output_folder, project_name = "samples", "sdxs"
    latents = generate_latents(
        embeddings=embeddings,
        height=height,
        width=width,
        num_inference_steps = num_inference_steps
    )

    images = decode_latents(latents, vae)

    os.makedirs(output_folder, exist_ok=True)
    for idx, image in enumerate(images):
        image.save(f"{output_folder}/{project_name}_{idx}.jpg")

    print("Images generated and saved to:", output_folder)

📚 ドキュメント

✨ 主な機能

Fast, Lightweight & Multilingual Diffusion for Everyone

私たちはAiArtLabという、予算が限られた小さな愛好家チームです。私たちの目標は、消費者向けグラフィックカードでトレーニングできる（LoRAではなく完全なトレーニングサイクル）コンパクトで高速なモデルを作成することです。小さなデータセットを効率的に処理し、16GBのGPU（例：RTX 4080）でも迅速にトレーニングできるU-Netを選択しました。私たちの予算は数千ドルに限られており、SDXLのような競合モデル（数千万ドル）よりも大幅に少ないため、SD1.5のような小さいが効率的なモデルを2025年向けに作成することにしました。

🔧 技術詳細

エンコーダアーキテクチャ（テキストと画像）

様々なエンコーダを実験した結果、LLaMAやT5 XXLのような大規模モデルは高品質な生成には必要ないと結論付けました。しかし、クエリのコンテキストを理解するエンコーダが必要で、「プロンプトの理解」に重点を置きました。80言語をサポートし、個々のトークンではなく文を処理する多言語エンコーダMexma-SigLIPを選択しました。Mexmaは最大512トークンを受け入れ、大きな行列を作成するため、トレーニングが遅くなります。そのため、プーリング層を使用して512x1152の行列を1x1152のベクトルに簡素化しました。具体的には、線形モデル/テキストプロジェクタを通してSigLIP埋め込みとの互換性を実現しました。これにより、テキスト埋め込みと画像を同期させることができ、統一されたマルチモーダルモデルにつながる可能性があります。この機能により、クエリで画像埋め込みとテキスト記述を混合することができます。さらに、モデルはテキスト記述なしで画像のみを使用してトレーニングすることができます。これは、アノテーションが難しいビデオのトレーニングを簡素化し、前のフレームの埋め込みを減衰させて入力することで、より一貫性のあるシームレスなビデオ生成を実現するはずです。将来的には、モデルを3D/ビデオ生成に拡張することを目指しています。

U-Netアーキテクチャ

滑らかなチャネルピラミッド[384, 576, 768, 960]を選択し、ブロックごとに2層、[4, 6, 8, 10]のトランスフォーマーを使用し、1152/48 = 24のアテンションヘッドを持ちます。このアーキテクチャは、約20億のパラメータを持つモデルで最高のトレーニング速度を提供し（私のRTX 4080に完全に収まります）、「深さ」が大きいため、「サイズ」が小さくてもSDXLと同等の品質が得られると考えています。1152層を追加することで、モデルを40億のパラメータに拡張でき、埋め込みサイズと完全に対称になり、エレガントさが高まり、おそらく「Flux/MJレベル」の品質が得られます。

VAEアーキテクチャ

SD3/Fluxの特徴である「ヘイズ」がなく、詳細、テキスト、解剖学的構造を保持する、従来とは異なる8x 16チャネルのAuraDiffusion VAEを選択しました。FFN畳み込みを使用した高速バージョンを使用しましたが、細かいパターンでわずかなテクスチャの損傷が見られ、ベンチマークでの評価が下がる可能性があります。ESRGANのようなアップスケーラーでこれらのアーティファクトを解消できます。全体的に、このVAEは非常に過小評価されていると考えています。

トレーニングプロセス

オプティマイザ

いくつかのオプティマイザ（AdamW、Laion、Optimi-AdamW、Adafactor、AdamW-8bit）をテストし、AdamW-8bitを選択しました。Optimi-AdamWは最も滑らかな勾配減衰曲線を示しましたが、AdamW-8bitはより混乱した動作を示します。しかし、その小さなサイズにより、より大きなバッチサイズが可能で、低コストのGPUでのトレーニング速度を最大化します（トレーニングには4台のA6000と5台のL40を使用しました）。

学習率

減衰/ウォームアップ曲線を操作することには効果がありますが、それほど大きくはありません。最適な学習率はしばしば過大評価されています。私たちの実験では、Adamは広い学習率範囲を許容します。トレーニング中に1e-4から徐々に1e-6に減少させました。言い換えると、正しいモデルアーキテクチャを選択することは、ハイパーパラメータを微調整するよりもはるかに重要です。

データセット

モデルは約100万枚の画像でトレーニングされました。256解像度のImageNetで60エポック（低品質のアノテーションのため時間を無駄にしました）、CaptionEmporium/midjourney-niji-1m-llavanextで8エポック、576解像度のリアルな写真とアニメ/アートを使用しました。人間のプロンプト、Caption Emporiumが提供するプロンプト、SmilingWolfのWD-Tagger、Moondream2を使用してアノテーションを行い、プロンプトの長さと構成を変えて、モデルがさまざまなプロンプトスタイルを理解できるようにしました。データセットは非常に小さいため、多くのエンティティを見逃し、「自転車に乗ったガチョウ」のような未知の概念に対応するのが難しいモデルになっています。データセットには多くのワイフスタイルの画像も含まれており、人間の解剖学的構造を学習する能力に興味があり、「馬に乗った宇宙飛行士」のような描画能力ではなく、モデルが人間の解剖学的構造を学習する能力に興味がありました。ほとんどの説明は英語でしたが、テストによるとモデルは多言語対応しています。

制限事項

• 非常に小さいデータセットのため、概念のカバレッジが限られています。
• Image2Image機能はさらなるトレーニングが必要です（テキストから画像へのトレーニングに焦点を当てるため、SigLIPの割合を5%に減らしました）。

謝辞

• Stan — 主要な投資家。主要な資金的支援 - 他の人が狂気だと言ったときに私たちを信じてくれてありがとう。
• Captainsaturnus — 物質的支援。
• Lovescape & Whargarbl — 精神的支援。
• CaptionEmporium — データセット。

"私たちは、効率的でコンパクトなモデルが未来を担うと信じています。寄付いただきありがとうございます。今後もご支援いただけることを期待しています。"

トレーニング予算

寄付

トレーニングにGPUや資金を提供できる場合は、ご連絡ください。

DOGE: DEw2DR8C7BnF8GgcrfTzUjSnGkuMeJhg83 BTC: 3JHv9Hb8kEW8zMAccdgCdZGfrHeMhH1rpN

連絡先

recoilme

📄 ライセンス

Apache-2.0