HiDream-O1-Image: Eigenes LoRA-Training trotz neuartiger Architektur

Die Landschaft der KI-gestützten Bildgenerierung entwickelt sich rasant weiter. HiDream-O1-Image, ein neues Open-Weights-Modell, sticht dabei besonders hervor. Es erreichte Platz 8 in der künstlichen Intelligenz-Bewertungsarena für Text-zu-Bild-Modelle – trotz seines ungewöhnlichen Ansatzes.

Doch während viele Modelle auf etablierten Architekturen wie SDXL oder Flux basieren, setzt HiDream-O1-Image auf eine komplett neuartige Struktur. Statt auf klassische Komponenten wie VAE oder separate Textencoder setzt das Modell auf einen einheitlichen Transformer-Ansatz. Das wirft eine zentrale Frage auf: Wie lässt sich ein LoRA für ein solches Modell trainieren, wenn die üblichen Tools wie kohya, ai-toolkit oder SimpleTuner nicht kompatibel sind?

Warum herkömmliche LoRA-Tools versagen

Die meisten LoRA-Trainer sind auf die gängige Architektur von Text-zu-Bild-Modellen ausgelegt. Sie gehen von einem UNet oder DiT-Denoiser aus, der mit einem VAE und einem oder mehreren Textencodern kombiniert wird. Die Trainingsprozesse sind darauf ausgelegt, diese Module zu patchen und anzupassen.

HiDream-O1-Image bricht mit diesem Paradigma. Stattdessen handelt es sich um ein Pixel-level Unified Transformer (UiT) – ein Modell, das keine externen VAE oder Textencoder benötigt. Stattdessen verarbeitet es Rohpixel, Text und Aufgabenbedingungen in einem einzigen, gemeinsamen Token-Raum. Der Kern des Modells ist ein Qwen3VLForConditionalGeneration-Transfomer, der direkt aus der Hugging Face-Bibliothek stammt.

Ein zentrales Merkmal ist die Patchifizierung: Ein Bild wird in 32x32-Pixel-Blöcke unterteilt, die zu Tokens umgewandelt werden. Jeder Token besteht aus einem 3072-dimensionalen Vektor, der die Rohpixelwerte repräsentiert. Es gibt keine separate latente Darstellung – das Modell arbeitet direkt auf den Pixeln.

Diese Architektur macht herkömmliche LoRA-Trainer unbrauchbar. Die Tools erwarten eine UNet-Struktur, in die sie LoRA-Adapter einfügen können. Doch HiDream-O1-Image hat weder ein UNet noch einen VAE. Stattdessen ist es ein reiner Transformer, der mit einem Flow-Matching-Ansatz trainiert wird.

Reverse-Engineering des Trainingsschemas

Da HiDream-O1-Image nur mit Inferenz-Code ausgeliefert wurde, musste der Trainingsprozess aus den vorhandenen Code-Snippets abgeleitet werden. Der Schlüssel lag im Verständnis der Inferenzschleife, die in der Datei models/pipeline.py dokumentiert ist.

Der Prozess beginnt mit der Erzeugung von Rauschen, das mit einem Skalierungsfaktor von 8.0 multipliziert wird. Dies entspricht der Gleichung:

z_t = (1 - σ) * x0 + σ * (8.0 * ε)

Dabei ist σ der Rauschpegel, x0 das saubere Bild und ε das zufällige Rauschen. Das Modell erhält diesen Wert z_t zusammen mit einem Timestep, der als 1.0 - σ berechnet wird. Die Ausgabe des Modells ist eine Vorhersage des sauberen Bildes x_pred.

Diese Erkenntnis führte zu einem entscheidenden Schluss: Das Modell ist darauf trainiert, das saubere Bild x0 vorherzusagen, nicht das Rauschen ε. Das ist ein fundamentaler Unterschied zu den meisten herkömmlichen Modellen, die auf Rauschen trainiert werden.

Die Loss-Funktion ist daher einfach ein MSE zwischen der Modellvorhersage `x_pred` und dem sauberen Bild `x0`, beschränkt auf die Bild-Token-Positionen.

Der erste öffentliche LoRA-Trainingsansatz

Auf Basis dieser Erkenntnisse wurde ein 150-Zeilen-Trainer entwickelt, der speziell für HiDream-O1-Image optimiert ist. Der Trainer nutzt die PEFT-Bibliothek, um LoRA-Adapter in die Linear-Layer des Modells einzufügen. Da das Backbone ein reiner Transformer ist, ist die Integration der Adapter unkompliziert.

Der Trainer funktioniert in folgenden Schritten:

• Datenvorbereitung: Die Trainingsbilder werden in 32x32-Pixel-Blöcke unterteilt und zu Tokens umgewandelt.
• Rauschgenerierung: Für jedes Bild wird ein zufälliger Rauschpegel σ erzeugt, und das Rauschen wird mit dem Faktor 8.0 skaliert.
• Modellvorhersage: Das Modell erhält das verrauschte Bild z_t und den Timestep 1.0 - σ und gibt die Vorhersage x_pred aus.
• Loss-Berechnung: Der MSE zwischen x_pred und dem sauberen Bild x0 wird berechnet und zur Optimierung genutzt.
• LoRA-Integration: Die LoRA-Adapter werden in die Linear-Layer des Modells eingefügt und trainiert.

Ein zentraler Punkt ist die Timestep-Behandlung. Während herkömmliche Modelle den Timestep oft als separaten Eingabewert behandeln, wird er bei HiDream-O1-Image als spezielles Token <|tms_token|> in den Eingabestrom integriert. Der Trainer muss diesen Mechanismus berücksichtigen.

Herausforderungen und Fallstricke

Die Entwicklung des Trainers war kein geradliniger Prozess. Mehrere Hürden mussten überwunden werden:

• Timestep-Berechnung: Die Umwandlung des Rauschpegels σ in den Timestep 1.0 - σ war nicht intuitiv. Eine falsche Berechnung führte zu inkonsistenten Ergebnissen.
• Token-Typen: Das Modell unterscheidet zwischen Bild- und Text-Tokens. Die Maskierung der Bild-Token-Positionen im Loss war entscheidend, um eine korrekte Berechnung zu gewährleisten.
• Rauschskalierung: Der Faktor 8.0 für das Rauschen war ein kritischer Parameter. Eine falsche Skalierung führte zu instabilen Trainingsprozessen.
• Adapter-Integration: Die korrekte Platzierung der LoRA-Adapter in den Linear-Layern erforderte ein tiefes Verständnis der Modellarchitektur.

Ein besonders kniffliger Punkt war die Verarbeitung der Bild-Tokens. Da das Modell direkt auf Pixeln arbeitet, mussten die Eingabedaten sorgfältig vorbereitet werden, um sicherzustellen, dass die Tokenisierung korrekt funktioniert.

Erste Ergebnisse und Anwendungsmöglichkeiten

Der entwickelte LoRA-Trainer produziert visuelle Verbesserungen, die in zwei Hauptkategorien fallen:

• Anime-Stil: Die Verbesserung der Bildqualität, Beleuchtung und Stilisierung bei Anime-ähnlichen Darstellungen.
• Halb-realistische Bilder: Die Optimierung von Licht, Kontrasten und Details bei halb-realistischen Motiven.

Die LoRA ist kein Charakter-LoRA, sondern ein allgemeiner Stilverbesserer. Sie funktioniert mit einem Trigger-Phrase, der beim Prompting verwendet wird.

Die ersten Tests zeigen vielversprechende Ergebnisse. Die LoRA verbessert die Bildqualität deutlich, ohne den ursprünglichen Stil des Modells zu verlieren.

Zukunftsperspektiven für HiDream-O1-Image

HiDream-O1-Image markiert einen Wendepunkt in der Entwicklung von Text-zu-Bild-Modellen. Die einheitliche Transformer-Architektur bietet neue Möglichkeiten, erfordert aber auch innovative Ansätze für die Anpassung.

Während dieser Beitrag einen der ersten öffentlichen LoRA-Trainingsansätze für das Modell vorstellt, bleibt noch viel Raum für weitere Entwicklungen.

• Erweiterung der LoRA-Funktionalität: Die Integration von LoRAs für spezifische Stile oder Charaktere könnte die Anwendungsmöglichkeiten deutlich erweitern.
• Optimierung des Trainingsprozesses: Durch Anpassung der Hyperparameter und Datenvorbereitung könnte die Qualität der LoRAs weiter verbessert werden.
• Community-Beiträge: Die Veröffentlichung des Trainers und der LoRA-Dateien soll die Community dazu ermutigen, eigene Experimente durchzuführen und das Modell weiterzuentwickeln.

Die Entwicklung von HiDream-O1-Image und die Möglichkeit, LoRAs für das Modell zu trainieren, öffnet die Tür für eine neue Ära der Bildgenerierung. Mit den richtigen Tools und einem tiefen Verständnis der Architektur können Entwickler und Enthusiasten das volle Potenzial dieses innovativen Ansatzes ausschöpfen.