Sign In

How to Train a Krea 2 LoRA on a 12GB GPU

0

How to Train a Krea 2 LoRA on a 12GB GPU

I wanted to find out whether a modern 12GB consumer GPU could train a useful Krea 2 LoRA locally without relying on cloud hardware.

The answer is yes—but only with aggressive memory offloading and the right training workflow.

This guide summarizes the setup I used with an RTX 3060 12GB, Musubi Tuner, Krea 2 Raw for training, and Krea 2 Turbo for fast inference.

Important workflow rule: train the LoRA on Krea 2 Raw, then apply the resulting weights to Krea 2 Turbo in ComfyUI.

Krea 2 responds well to detailed textual prompts and can also interpret JSON- or BOX-like formatting as a textual prompting convention. Unlike Ideogram 4, however, it does not expose a documented native JSON schema or deterministic bounding-box control interface.

Krea 2 Raw is the flexible base checkpoint intended for fine-tuning, while Turbo is the distilled 8-step model intended for fast generation.


What made this difficult on 12GB VRAM?

The full training stack is much larger than the available GPU memory:

  • Krea 2 Raw: approximately 24.48GB

  • Qwen3-VL 4B text encoder: approximately 8.27GB

  • Qwen Image VAE: approximately 1.1GB

Loading everything in BF16 or FP16 immediately exceeds the capacity of a 12GB card. Using FP8 alone was not enough in my tests: training remained unstable, slow, and prone to CUDA crashes.

The stable solution was FP8 base-model quantization combined with CPU block swapping.


The low-VRAM configuration

These four options made the largest difference:

  • --fp8_base --fp8_scaled reduces the base model’s training memory footprint.

  • --blocks_to_swap 20 moves inactive transformer blocks to system RAM.

  • --block_swap_h2d_only keeps a CPU master copy of the frozen base weights and streams blocks from host to device without copying them back. This removes an unnecessary device-to-host transfer during LoRA training.

  • --block_swap_ring_size 2 uses two GPU ring buffers so the next block can be prefetched while the current block is processed. This can reduce transfer-related GPU idle periods. A ring size of 2 is also the current Musubi default, but it is specified explicitly for reproducibility.

Attribution note: the benchmark tests the combined block-swapping configuration. It does not isolate --block_swap_h2d_only from the other settings, so no independent percentage speedup is claimed for that option.

System RAM matters. I tested this workflow with 64GB of RAM. It may work with less, but 32GB can become tight depending on the dataset, Windows background usage, caching, and page-file settings.


My test hardware and training target

  • GPU: NVIDIA GeForce RTX 3060 12GB

  • System RAM: 64GB

  • Operating system: Windows

  • Validated bucket geometries: 512 × 512 and 1024 × 256

  • Batch size: 1

  • LoRA rank / alpha: 16 / 16

  • Optimizer: AdamW 8-bit

  • Learning rate: 1e-4


1. Install Musubi Tuner

Install Git, Python 3.10 or 3.11, a recent NVIDIA driver, and optionally the Microsoft Visual Studio Build Tools.

cd C:\
git clone https://github.com/kohya-ss/musubi-tuner.git C:\musubi-tuner
cd C:\musubi-tuner

python -m venv .venv
.\.venv\Scripts\activate

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
pip install -e .
pip install ascii-magic matplotlib tensorboard prompt-toolkit

The PyTorch command above matches the CUDA 12.4 environment I tested. Use a PyTorch build appropriate for your own driver and environment.


2. Prepare the model files

I used the following folder structure:

C:\krea2_models\
├─ DiffusionModels\
│  └─ raw.safetensors
├─ text_encoders\
│  └─ qwen3vl_4b_bf16.safetensors
└─ VAE\
   └─ qwen_image_vae.safetensors

License note: Krea 2 uses the Krea 2 Community License. Check the current terms before distributing or commercially using trained LoRAs and generated outputs.


3. Prepare the dataset

Each training image should have a matching .txt caption file:

C:\krea2_lora_training\mp_concept_krea2\
├─ krea2_images\
│  ├─ image_0001.png
│  ├─ image_0001.txt
│  ├─ image_0002.png
│  ├─ image_0002.txt
│  └─ ...
├─ krea2_cache_512\
└─ output_krea2\

I reused structured JSON captions because they are easy to audit and can organize object, style, and layout descriptions consistently. This is my own dataset convention rather than a strict Krea 2 requirement. Krea 2 receives the serialized JSON as prompt text, so BOX- or bbox-like fields should be treated as textual formatting cues rather than guaranteed spatial constraints. Natural-language captions remain a valid and often simpler alternative. See the Krea 2 Technical Report for the broader prompting context.

Example caption:

{"high_level_description":"A clean game-art reference image of a futuristic off-road utility buggy, mp_concept_krea2.","style_description":{"aesthetics":"production-friendly vehicle concept art, readable silhouette, game asset clarity","lighting":"soft studio lighting with clear material readability","medium":"digital vehicle concept art","art_style":"hard-surface game vehicle design with large wheels, protective roll cage, modular panels, elevated suspension, and clean forms"},"compositional_deconstruction":{"background":"A simple neutral studio background with subtle shading.","elements":[{"type":"obj","desc":"A single futuristic off-road utility buggy shown in a clean three-quarter view, with large wheels, exposed frame, protective roll cage, compact chassis, and readable game-asset proportions, mp_concept_krea2."}]}}

Keep the trigger token consistent across the dataset and remove irrelevant automatic tags or metadata.


4. Create the dataset configuration

Save this as mp_concept_krea2_dataset_512.toml:

[general]
resolution = [512, 512]
caption_extension = ".txt"
batch_size = 1
enable_bucket = true
bucket_no_upscale = false

[[datasets]]
image_directory = "C:/krea2_lora_training/mp_concept_krea2/krea2_images"
cache_directory = "C:/krea2_lora_training/mp_concept_krea2/krea2_cache_512"
num_repeats = 1

5. Cache the dataset

Run the following commands in PowerShell. Cache the VAE latents first:

cd C:\musubi-tuner

.\.venv\Scripts\python.exe src\musubi_tuner\krea2_cache_latents.py `
  --dataset_config C:\krea2_lora_training\mp_concept_krea2\mp_concept_krea2_dataset_512.toml `
  --vae C:\krea2_models\VAE\qwen_image_vae.safetensors

Then cache the text-encoder outputs:

.\.venv\Scripts\python.exe src\musubi_tuner\krea2_cache_text_encoder_outputs.py `
  --dataset_config C:\krea2_lora_training\mp_concept_krea2\mp_concept_krea2_dataset_512.toml `
  --text_encoder C:\krea2_models\text_encoders\qwen3vl_4b_bf16.safetensors `
  --batch_size 1

6. Launch the optimized training run

cd C:\musubi-tuner

.\.venv\Scripts\accelerate.exe launch --num_cpu_threads_per_process 1 --mixed_precision bf16 `
  src\musubi_tuner\krea2_train_network.py `
  --dit C:\krea2_models\DiffusionModels\raw.safetensors `
  --fp8_base `
  --fp8_scaled `
  --vae C:\krea2_models\VAE\qwen_image_vae.safetensors `
  --dataset_config C:\krea2_lora_training\mp_concept_krea2\mp_concept_krea2_dataset_512.toml `
  --sdpa `
  --mixed_precision bf16 `
  --timestep_sampling krea2_shift `
  --weighting_scheme none `
  --optimizer_type adamw8bit `
  --learning_rate 1e-4 `
  --gradient_checkpointing `
  --blocks_to_swap 20 `
  --block_swap_h2d_only `
  --block_swap_ring_size 2 `
  --max_data_loader_n_workers 2 `
  --persistent_data_loader_workers `
  --network_module networks.lora_krea2 `
  --network_dim 16 `
  --network_alpha 16 `
  --max_train_steps 500 `
  --save_every_n_steps 50 `
  --save_state `
  --save_precision bf16 `
  --seed 42 `
  --output_dir C:\krea2_lora_training\mp_concept_krea2\output_krea2 `
  --output_name mp_concept_krea2_lora_512

The launcher uses krea2_shift, which applies Krea 2’s resolution-aware timestep schedule independently to each sample and bucket. The previous fixed value of 2.5 approximately matched 1024 × 1024 inference and was therefore not ideal for a 512-pixel-area, aspect-ratio-bucketed dataset. Do not combine krea2_shift with --discrete_flow_shift 2.5.

This is a scheduler correction, not an acceleration setting. The speed improvement comes primarily from keeping the FP8 and block-swapping path active.

Why rank 16 / alpha 16?

Rank 16 / alpha 16 is a deliberate middle ground based on the practical comparisons in my Ideogram 4 LoRA training guide. In those tests, rank 8 produced a lighter style influence that was easier to stack with other LoRAs, rank 16 provided a useful balance between concept strength and composability, and rank 32 was better suited to a precise object or concept intended to take priority over other stacked LoRAs.

These observations are workflow-specific rather than an official Krea 2 rule. Musubi currently documents rank 32 / alpha 32 as the model-author default, while this guide intentionally starts at 16 / 16. Rank 16 alone is not what makes 12GB training possible: FP8 quantization and block swapping are far more important to memory feasibility.


Performance on the RTX 3060 12GB

The validated reference configuration combines FP8 scaled base weights, 20 swapped transformer blocks, host-to-device-only streaming, a two-buffer block-swap ring, gradient checkpointing, SDPA, batch size 1, rank 16 / alpha 16, and resolution-aware krea2_shift timestep sampling.

The successive tests on the reference RTX 3060 12GB system show how much the complete optimization path matters:

  • Krea 2 Raw BF16 — no block swap
    Stability: CUDA OOM
    Observed speed: ~154 s/step before the crash

  • Krea 2 Raw FP8 — no block swap
    Stability: Unstable, CUDA crash
    Observed speed: ~85–113 s/step

  • Krea 2 Raw FP8 — validated optimized configuration
    Stability: Stable
    Average speed: ~7.8 s/step

The validated run maintained 100% GPU utilization and used approximately 11.94GB of VRAM. Similar performance across square and horizontal image sets confirms that the speedup comes from the optimized FP8 and block-swapping pipeline, not from a particular bucket geometry.

Block swapping was required on this tested 12GB setup. The benchmark validates the combined configuration; it does not isolate H2D-only streaming or ring size from the other settings, so no independent percentage speedup is attributed to either option.


Inference in ComfyUI

After training, copy the resulting .safetensors file into:

ComfyUI\models\loras

For validation and production inference:

  1. Load krea2_turbo_fp8_scaled.safetensors as the base diffusion model.

  2. Apply the trained LoRA with a standard LoRA loader on the model path.

  3. Start with a LoRA strength between 0.8 and 1.2.

  4. Include the trigger token in the prompt exactly as it appeared in the training captions.

A JSON/BOX builder such as KJNodes can preserve the same auditable text convention and place the trigger token consistently inside descriptions. Krea 2 still receives the serialized result as prompt text: unlike Ideogram 4, BOX or bbox fields are textual formatting cues, not documented deterministic spatial controls.

vader_turbo_8_52003.png

What I would change for a larger training run

  • Validate several checkpoints before committing to thousands of steps.

  • Choose rank deliberately: start with 8 for a lighter, stackable style influence; use 16 as a balanced default; test 32 when a precise object or dominant concept needs more capacity.

  • Keep the first experiments at 512px until the dataset and captions are proven.

  • Monitor system RAM and page-file use, not only VRAM.

  • Test the LoRA on Turbo regularly because that is the intended inference target.


Final thoughts

Krea 2 LoRA training is possible on a 12GB GPU, but FP8 alone is not enough.

The practical breakthrough came from combining FP8 base-model loading with transformer block swapping, one-way host-to-device transfers, double buffering, gradient checkpointing, and pre-cached dataset features.

On the reference system, the validated configuration reached approximately 7.8 s/step with both square and horizontal buckets. Results on other hardware and resolutions may differ, but the matching geometries show that the optimized pipeline—not the horizontal format—accounts for the performance gain. The workflow is most relevant when Krea 2 Turbo is the intended final inference model.

Krea 2 support in Musubi Tuner is still experimental, so commands and recommended settings may change as the project evolves.


0