Stable Diffusion 3.5 von Grund auf neu in reinem PyTorch implementiert

(github.com/yousef-rafat)

2 Punkte von GN⁺ 2025-06-15 | 1 Kommentare | Auf WhatsApp teilen

Das Projekt miniDiffusion ist ein Open-Source-Projekt, das das Modell Stable Diffusion 3.5 von Grund auf ausschließlich mit PyTorch neu implementiert
Die Struktur dieses Projekts zeichnet sich dadurch aus, dass sie auf Lernzwecke sowie Experimente und Hacking ausgerichtet ist
Die gesamte Codebasis umfasst etwa 2800 Zeilen und besteht mit minimalem Code aus Komponenten vom VAE über DiT bis hin zu Trainings- und Datensatz-Skripten
Zu den wichtigsten Bestandteilen gehören VAE, CLIP, T5-Text-Encoder, ein multimodaler Diffusion Transformer sowie Joint Attention
Es enthält noch experimentelle Funktionen und benötigt daher weitere Tests

Vorstellung des miniDiffusion-Projekts

miniDiffusion ist ein Open-Source-Projekt, das die Kernfunktionen von Stable Diffusion 3.5 ausschließlich mit PyTorch neu implementiert.
Im Vergleich zu bestehendem Stable Diffusion 3.5 bietet dieses Projekt unter anderem folgende Vorteile:

Mit einer Codebasis von rund 2.800 Zeilen ist es sehr kompakt und daher hervorragend geeignet, um die Struktur direkt zu analysieren und zu erlernen
Es lässt sich nützlich für verschiedenste Machine-Learning-Experimente und Model-Hacking einsetzen
Es hat nur sehr wenige Abhängigkeiten und nutzt nur ein Minimum an Bibliotheken

Kernstruktur und zentrale Dateien

dit.py : Implementierung des Hauptmodells von Stable Diffusion
dit_components.py : Komponenten für Embeddings, Normalisierung, Patch-Embedding und Hilfsfunktionen für DiT
attention.py : Implementierung des Joint-Attention-Algorithmus
noise.py : enthält einen Euler-ODE-Scheduler für Rectified Flow
t5_encoder.py, clip.py : Implementierung der T5- und CLIP-Text-Encoder
tokenizer.py : Implementierung von Byte-Pair- und Unigram-Tokenizern
metrics.py : Implementierung der Bewertungsmetrik FID (Fréchet inception distance)
common.py : stellt Hilfsfunktionen für das Training bereit
common_ds.py : Implementierung eines iterierbaren Datensatzes, der Bilder in Trainingsdaten für DiT umwandelt
model-Ordner : speichert Modell-Checkpoints und Logs nach dem Training
encoders-Ordner : speichert Checkpoints separater Module wie VAE und CLIP

⚠️ Experimentelle Funktionen und Testbedarf miniDiffusion enthält noch experimentelle Funktionen und benötigt weitere Tests

Detaillierte Zusammensetzung nach Hauptfunktionen

Core Image Generation Modules

Implementierung von VAE, CLIP, T5-Text-Encodern
Implementierung von Byte-Pair- und Unigram-Tokenizern

SD3 Components

Multi-Modal Diffusion Transformer Model
Implementierung eines Flow-Matching Euler Scheduler
Logit-Normal Sampling
Einführung des Joint-Attention-Algorithmus

Skripte für Training und Inferenz

Bereitstellung von Trainings- und Inferenz-Skripten für SD3 (Stable Diffusion 3.5)

Lizenz

Veröffentlicht unter der MIT-Lizenz und erstellt für Lern- und Experimentierzwecke

Bedeutung und Vorteile dieses Open-Source-Projekts

Eine aktuelle Bildgenerierungsmodell-Architektur auf dem Niveau von Stable Diffusion 3.5 lässt sich direkt in reinem PyTorch trainieren und hacken
Der Code ist kompakt und unabhängig und damit optimal für Strukturanalyse/Modell-Tuning/Forschung an neuen Algorithmen
Aktuelle multimodale, Transformer- und Attention-Verfahren lassen sich direkt praktisch erproben
Bietet eine Grundlage für sicheres Experimentieren unabhängig von kommerziellen Projekten

1 Kommentare

GN⁺ 2025-06-15

Hacker-News-Kommentare

Die Flux-Referenzimplementierung ist wirklich sehr minimal aufgebaut, daher lohnt sich ein Blick für alle, die daran interessiert sind
- Flux GitHub
- Das minRF-Projekt hat den Vorteil, dass man mit rectified flow leicht in das Training kleiner Diffusionsmodelle einsteigen kann
- minRF GitHub
- Auch die Referenzimplementierung von Stable Diffusion 3.5 ist ziemlich kompakt geschrieben und daher gut als Vorlage geeignet
- SD 3.5 GitHub
- Referenzimplementierungen werden oft nicht gut gepflegt und enthalten viele Bugs
  - Als Beispiel kann man den CLIP-Tokenizer von OpenAI nennen: Die inoffizielle Referenzversion, nicht die für das offizielle Training verwendete Version, wird trotz Bugs unverändert in viele Projekte kopiert
  - Auch Flux wurde nicht tatsächlich für das Training verwendet und hat Bugs, die unter anderem bei cudagraphs kleinere Probleme verursachen können
  - Es hängt von der CLIP-Referenz ab, und da CLIP selbst Bugs hat, werden diese entsprechend weitervererbt
Es stellt sich die Frage, ob das miniDiffusion-Projekt bedeutet, dass es das Stable-Diffusion-3.5-Modell verwendet
- Zugehöriger Code
- Der Trainingsdatensatz ist sehr klein und enthält nur modebezogene Bilder
- Fashion-Datensatz
- Dieser Datensatz dient dazu, das Fine-Tuning eines Diffusionsmodells praktisch auszuprobieren
  - Der Code ist eine Neuimplementierung von SD3, aber wegen der Grenzen meiner Hardware wurden die Gewichte von HuggingFace übernommen
Es stellt sich die Frage, ob die Verwendung von reinem PyTorch Leistungsvorteile auf GPUs bringt, die nicht von NVIDIA sind, oder ob PyTorch so stark auf CUDA optimiert ist, dass andere GPU-Anbieter nicht konkurrenzfähig sein können
- PyTorch funktioniert auch auf Apple Silicon ziemlich gut
  - Allerdings ist ein direkter Vergleich schwierig, da Apple-GPUs rechnerisch kaum die Leistung von NVIDIAs Spitzenmodellen erreichen
  - Der Vollständigkeit halber gibt es bei der Nutzung von PyTorch auf Apple Silicon eine kleine Besonderheit
    - Jeder Tensor gilt als einem bestimmten Gerät (CPU oder GPU) „zugehörig“, daher wird beim Verschieben der Daten immer eine vollständige Kopie erzeugt
    - Macs haben zwar Unified Memory, aber PyTorch ist weiterhin so aufgebaut, dass es Datenkopien durchführt
- Auch auf Nicht-NVIDIA-Geräten wie AMD kann man ML-Workloads über Vulkan ausführen
  - Wenn Erweiterungen wie cooperative matrix und neue Funktionen auf Treiberebene weiter unterstützt werden, könnte selbst ein Leistungsunterschied von nur einigen Prozent gegenüber CUDA künftig verschwinden
- Die ROCm-Unterstützung in PyTorch kommt nur sehr langsam voran, und selbst wenn sie funktioniert, ist sie langsam
- PyTorch funktioniert unter ROCm zwar gut, aber ob es wirklich vollständig „auf Augenhöhe“ ist, weiß ich nicht

Es gibt den Vorschlag, im PyTorch-Code statt

    self.q = nn.Linear(embed_size, embed_size, bias = False)
    self.k = nn.Linear(embed_size, embed_size, bias = False)
    self.v = nn.Linear(embed_size, embed_size, bias = False)

lieber

    self.qkv = nn.Linear(embed_size, 3*embed_size, bias = False)
    # innerhalb der forward-Funktion
    qkv = self.qkv(x)

zu versuchen

Dadurch wären die Parameter von q, k und v nicht mehr jeweils getrennt verbunden wie zuvor, sondern zwischen q, k und v gekoppelt
- Falls ich gerade nur zu müde bin und etwas verwechsle, bitte ich um Nachsicht
Das wirkt wie gutes Material für Lernende
- Es stellt sich die Frage, ob es auch ein Tutorial oder eine Anleitung gibt, der selbst Anfänger folgen können
- Bei fast.ai gibt es einen Kurs, in dem Stable Diffusion direkt implementiert wird
  - fast.ai-Kurs
Es stellt sich die Frage, ob das bedeutet, dass man Stable Diffusion ohne Lizenzbeschränkungen verwenden kann
- Nein, die Inferenz-/Trainingsalgorithmen (also die Mathematik selbst) unterliegen nicht dem Urheberrecht, aber der OP hat lediglich den Code neu implementiert
  - Urheberrechtlich relevant ist das Modell selbst, also die Gewichte, und der OP hat mangels Daten und Rechenleistung nicht selbst trainiert
Ehrlich gesagt frage ich mich, vielleicht etwas peinlich, was wir durch dieses Repo im Vergleich zu vorher eigentlich neu gewonnen haben
- Ich habe persönlich den Bau von Modellen eher gemieden und hauptsächlich nur die Ergebnisse aus der Distanz verfolgt
- Ich war vage davon ausgegangen, dass es bereits Inferenz-/Trainingsskripte auf PyTorch-Basis öffentlich gibt
- Zumindest dachte ich, dass Inferenzskripte bei der Auslieferung eines Modells ohnehin mitgeliefert würden, und dass es auch Fine-Tuning-/Trainingsskripte geben müsse
- Ich bin nicht sicher, ob dieses Projekt eine Art „Clean Room“- oder „Dirty Room“-Neuschreibung bestehender Dinge ist oder ob vorhandener PyTorch-Code selbst schon so stark mit CUDA/C verkompliziert ist, dass eine reine PyTorch-Version einen großen Mehrwert hat
- Jedenfalls kenne ich mich nicht gut genug aus und würde mich freuen, wenn das jemand erklären könnte
  - Der zentrale Wert dieses Projekts liegt in einer „Implementierung mit minimalen Abhängigkeiten“
    - Ich habe SD 3.5 selbst nie ausgeführt, aber es basiert auf der huggingface-Bibliothek, und persönlich empfinde ich die Abhängigkeiten von huggingface als so komplex, dass sich Modelle oft kaum starten lassen, wenn die Umgebung nicht fast identisch mit der des Entwicklers ist
    - Gerade einige Monate oder Jahre nach dem ursprünglichen Release wird das Ausführen bestimmter Modelle besonders schwierig
    - Wenn man sich zum Beispiel die Datei requirements.txt der Referenzimplementierung von Stability AI für SD3.5 ansieht, sind dort keine Versionen festgelegt, und es werden riesige Bibliotheken wie transformers eingebunden, was in der Praxis wirklich unerquicklich ist
  - Stability AI veröffentlicht die Stable-Diffusion-Modelle unter der Stability AI Community License, was anders als MIT nicht „völlig frei“ ist
    - Änderungen an den Gewichten auf bestimmte Weise sind nicht erlaubt
    - Dieses Paket ist so aufgebaut, dass Modellausführung (Inferenz) oder Fine-Tuning mit bereits vorhandenen AI-Gewichten möglich ist
    - Zum Lernen und Studieren ist es hervorragend geeignet, aber die Lizenzfragen bestehen weiterhin
Wenn ich an SD 3.5 denke, oder an jede andere Version, dann erscheint mir der entscheidende Teil in den während des Trainings erzeugten Gewichten zu liegen
- Der Code selbst ist im Hinblick auf Qualität oder Leistung des Ergebnisses vergleichsweise weniger wichtig
- Das ist aber kein endgültiges Urteil und soll diese Arbeit keineswegs herabsetzen
Ich frage mich, wie praxistauglich der ursprünglich vom CompViz-Team der Ludwig-Maximilians-Universität veröffentlichte akademische Quellcode ist
Es stellt sich die Frage, ob die Diffusion-Transformer-(DiT-)Implementierung hier Cross-Token-Attention so vollständig umsetzt wie die Vollversion von SD 3.5, oder ob sie für bessere Lesbarkeit des Codes vereinfacht wurde