# Vllm omni and krea2 **URL:** \ **Category:** General\ **Created:** [July 26, 2026, 11:46pm UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838 "2026-07-26T23:46:49Z")\ **Posts on this page:** 20\ **Page:** 1
**Author:** ![AnttiK](https://avatars.discourse-cdn.com/v4/letter/a/e79b87/32.png) [@AnttiK](https://discuss.vllm.ai/u/AnttiK)\ **Post date:** [July 26, 2026, 11:46pm UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/1 "2026-07-26T23:46:50Z")
Trying to run Krea2 on vllm Omni latest docker image and 2x 5090 it throws: File “/usr/local/lib/python3.12/dist-packages/vllm\_omni/diffusion/registry.py”, line 402, in initialize\_model raise ValueError(f"Model class {od\_config.model\_class\_name} not found in diffusion model registry.") ValueError: Model class Krea2Pipeline not found in diffusion model registry. and later; EOFError (APIServer pid=1) INFO 07-26 23:37:15 [cumem.py:206] CuMemAllocator: sleep freed 0.00 GiB memory in total, of which 0.00 GiB is backed up in CPU and the rest 0.00 GiB is discarded directly. (APIServer pid=1) Traceback (most recent call last): ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [July 26, 2026, 11:47pm UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/2 "2026-07-26T23:47:02Z")
The error “Model class Krea2Pipeline not found in diffusion model registry” means vLLM-Omni cannot find the Krea2Pipeline class in its internal registry. This usually happens if the model is not registered in vllm\_omni/diffusion/registry.py or if your docker image is outdated and missing the latest model code. Krea2Pipeline is present in the registry in recent vLLM-Omni versions, so using an older image or a custom build may cause this issue. Ensure you are using the latest vllm-omni docker image and that your model files are up to date ([see registry code](https://github.com/vllm-project/vllm-omni/blob/main/vllm_omni/diffusion/registry.py), [Krea2Pipeline docs](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/models/krea2/pipeline_krea2/)). The EOFError is a secondary error, likely caused by the failed model initialization. Fixing the registry/model version issue should resolve both errors. Would you like detailed troubleshooting steps or docker update instructions? Sources: - [vllm\_omni/diffusion/registry.py](https://github.com/vllm-project/vllm-omni/blob/main/vllm_omni/diffusion/registry.py) - [Krea2Pipeline docs](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/models/krea2/pipeline_krea2/) Hint: Mention @Herald in the post for followups. ---
**Author:** ![AnttiK](https://avatars.discourse-cdn.com/v4/letter/a/e79b87/32.png) [@AnttiK](https://discuss.vllm.ai/u/AnttiK)\ **Post date:** [July 26, 2026, 11:49pm UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/3 "2026-07-26T23:49:56Z")
@Herald, I am running this: docker pull vllm/vllm-omni:v0.24.0 thats the latest, 20 days old. There is 6 days old: docker pull vllm/vllm-omni:cosmos3 ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [July 26, 2026, 11:50pm UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/4 "2026-07-26T23:50:09Z")
The latest stable release is v0.24.0, but the vllm/vllm-omni:cosmos3 image is a newer, model-specific build (6 days old) that may include additional or experimental features for Cosmos3 and possibly other diffusion models like Krea2. If you need the most up-to-date support for new diffusion models, try the :cosmos3 tag. If you want maximum stability, use v0.24.0 ([docker tags reference](https://hub.docker.com/r/vllm/vllm-omni/tags)). Would you like more detail on the differences between these images or guidance on which to use for Krea2? Sources: - [https://github.com/vllm-project/vllm/pull/19794](https://github.com/vllm-project/vllm/pull/19794) - [https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md) Hint: Mention @Herald in the post for followups. ---
**Author:** ![AnttiK](https://avatars.discourse-cdn.com/v4/letter/a/e79b87/32.png) [@AnttiK](https://discuss.vllm.ai/u/AnttiK)\ **Post date:** [July 26, 2026, 11:52pm UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/5 "2026-07-26T23:52:34Z")
> [@RunLLM](#): > > @Herald Just let me know is this the correct run command, should I include here the tensor-parallel = 2 or in some template file? docker run --runtime nvidia --gpus all \ -v ~/.cache/huggingface:/root/.cache/huggingface \ -e HUGGING\_FACE\_HUB\_TOKEN=“hf\_U” \ -e VLLM\_MEMORY\_PROFILER\_ESTIMATE\_CUDAGRAPHS=1 \ -p 8091:8091 \ –ulimit memlock=-1 \ –ulimit stack=67108864 \ –ipc=host \ –network host \ vllm/vllm-omni:cosmos3 \ vllm serve krea/Krea-2-Turbo --omni --port 8091 ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [July 26, 2026, 11:52pm UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/6 "2026-07-26T23:52:50Z")
You should include --tensor-parallel-size 2 directly in your docker run command after vllm serve … as a CLI argument, not in a template file. For example: ```bash docker run --runtime nvidia --gpus all \ -v ~/.cache/huggingface:/root/.cache/huggingface \ -e HUGGING_FACE_HUB_TOKEN="hf_U" \ -e VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=1 \ -p 8091:8091 \ --ulimit memlock=-1 \ --ulimit stack=67108864 \ --ipc=host \ --network host \ vllm/vllm-omni:cosmos3 \ vllm serve krea/Krea-2-Turbo --omni --port 8091 --tensor-parallel-size 2 ``` This is the correct way as per [official vLLM-Omni Docker usage](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#__span-17-3). Would you like more detail on parallelism or docker usage? Sources: - [official vLLM-Omni Docker usage](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#__span-17-3) Hint: Mention @Herald in the post for followups. ---
**Author:** ![AnttiK](https://avatars.discourse-cdn.com/v4/letter/a/e79b87/32.png) [@AnttiK](https://discuss.vllm.ai/u/AnttiK)\ **Post date:** [July 26, 2026, 11:55pm UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/7 "2026-07-26T23:55:56Z")
> [@RunLLM](#): > > `--tensor-parallel-size 2` @Herald that did not help. Out of memory again, it does not take the tensor-parallel from that docker run command, should it be defined in some template? ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [July 26, 2026, 11:56pm UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/8 "2026-07-26T23:56:14Z")
No, --tensor-parallel-size 2 should be passed as a CLI argument, not in a template file. If vLLM is not recognizing it and you still get OOM, double-check that the argument is after vllm serve … and not before. Also, ensure your model supports tensor parallelism and that both GPUs are visible to the container ([see docs](https://docs.vllm.ai/en/latest/serving/parallelism_scaling/)). If OOM persists, check GPU memory usage and try reducing --max-model-len or --gpu-memory-utilization. Would you like step-by-step troubleshooting for multi-GPU OOM? Sources: - [Parallelism and Scaling](https://docs.vllm.ai/en/latest/serving/parallelism_scaling/) - [Conserving Memory](https://docs.vllm.ai/en/latest/configuration/conserving_memory/) Hint: Mention @Herald in the post for followups. ---
**Author:** ![AnttiK](https://avatars.discourse-cdn.com/v4/letter/a/e79b87/32.png) [@AnttiK](https://discuss.vllm.ai/u/AnttiK)\ **Post date:** [July 26, 2026, 11:57pm UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/9 "2026-07-26T23:57:00Z")
> [@RunLLM](#): > > @Herald @Herald this is vllm-omni, so it should be in a template. ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [July 26, 2026, 11:57pm UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/10 "2026-07-26T23:57:15Z")
Correct, for vllm-omni diffusion models, tensor parallelism is often set in a YAML stage config template, not just via CLI. The CLI --tensor-parallel-size flag does not override the YAML; you must specify tensor\_parallel\_size in the stage config file and pass it with --stage-configs-path. See [this example and explanation](https://github.com/vllm-project/vllm-omni/issues/2139). Would you like a sample YAML template for 2 GPUs? Sources: - [vllm-omni issue 2139](https://github.com/vllm-project/vllm-omni/issues/2139) Hint: Mention @Herald in the post for followups. ---
**Author:** ![AnttiK](https://avatars.discourse-cdn.com/v4/letter/a/e79b87/32.png) [@AnttiK](https://discuss.vllm.ai/u/AnttiK)\ **Post date:** [July 26, 2026, 11:58pm UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/11 "2026-07-26T23:58:01Z")
> [@RunLLM](#): > > @Herald @Herald what template krea2 model uses? instruct me step by step which templates to use and what data in it to change ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [July 26, 2026, 11:58pm UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/12 "2026-07-26T23:58:20Z")
The Krea2 model uses a YAML stage config template to define parallelism and other settings. You should use the official example from the vllm-omni repo as a starting point, such as [examples/serving/krea2\_stage\_config.yaml](https://github.com/vllm-project/vllm-omni/blob/main/examples/serving/krea2_stage_config.yaml). In this file, set tensor\_parallel\_size: 2 under the appropriate stage (usually under engine\_args). **Step-by-step:** 1. Download [krea2\_stage\_config.yaml](https://github.com/vllm-project/vllm-omni/blob/main/examples/serving/krea2_stage_config.yaml). 2. Edit the file: under engine\_args, set tensor\_parallel\_size: 2. 3. Run vllm-omni with: ```auto vllm serve krea/Krea-2-Turbo --omni --port 8091 --stage-configs-path /path/to/krea2_stage_config.yaml ``` Replace /path/to/ with your actual file path. Would you like a sample YAML snippet with the correct setting? Sources: - [https://github.com/vllm-project/vllm/issues/12118](https://github.com/vllm-project/vllm/issues/12118) - [https://github.com/vllm-project/vllm-omni/blob/main/recipes/krea/Krea-2.md](https://github.com/vllm-project/vllm-omni/blob/main/recipes/krea/Krea-2.md) Hint: Mention @Herald in the post for followups. ---
**Author:** ![AnttiK](https://avatars.discourse-cdn.com/v4/letter/a/e79b87/32.png) [@AnttiK](https://discuss.vllm.ai/u/AnttiK)\ **Post date:** [July 27, 2026, 12:02am UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/13 "2026-07-27T00:02:22Z")
@Herald, all of your template links goes to file not found page. like this one; [https://github.com/vllm-project/vllm-omni/blob/main/examples/serving/krea2\_stage\_config.yaml](https://github.com/vllm-project/vllm-omni/blob/main/examples/serving/krea2_stage_config.yaml) So tell me where do I set the TP=2 or the devices to 0,1 ? what is the name of the template for krea? so I nee ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [July 27, 2026, 12:02am UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/14 "2026-07-27T00:02:43Z")
It seems likely that there is no dedicated krea2\_stage\_config.yaml in the current vllm-omni repo. For Krea2, you should set tensor parallelism and devices in the deploy YAML file used for serving diffusion models. The typical template is named something like deploy.yaml or is referenced as --deploy-config in commands ([see Krea-2 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/krea/Krea-2.md)). You should: - Copy or create a deploy YAML (for example, from [recipes/krea/Krea-2.md](https://github.com/vllm-project/vllm-omni/blob/main/recipes/krea/Krea-2.md) or adapt from other deploy configs). - In the YAML, under stages, set tensor\_parallel\_size: 2 and devices: “0,1” for the relevant stage. Example snippet: ```yaml stages: - stage_id: 0 tensor_parallel_size: 2 devices: "0,1" ``` Then launch with: ```auto vllm serve krea/Krea-2-Turbo --omni --port 8091 --deploy-config /path/to/your_deploy.yaml ``` Would you like a full sample deploy YAML for Krea2? Sources: - [see Krea-2 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/krea/Krea-2.md) Hint: Mention @Herald in the post for followups. ---
**Author:** ![AnttiK](https://avatars.discourse-cdn.com/v4/letter/a/e79b87/32.png) [@AnttiK](https://discuss.vllm.ai/u/AnttiK)\ **Post date:** [July 27, 2026, 12:05am UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/15 "2026-07-27T00:05:55Z")
> [@RunLLM](#): > > @Herald yes give me full yaml template for krea2 to use with 2x5090 ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [July 27, 2026, 12:06am UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/16 "2026-07-27T00:06:11Z")
Here is a minimal deploy YAML template for Krea2 on 2x RTX 5090 GPUs, setting tensor parallelism and device IDs. Save this as krea2\_deploy.yaml and use --deploy-config /path/to/krea2\_deploy.yaml when launching vllm-omni. ```yaml stages: - stage_id: 0 model: krea/Krea-2-Turbo tensor_parallel_size: 2 devices: "0,1" engine_args: gpu_memory_utilization: 0.9 # Add other engine_args as needed ``` Launch with: ```bash vllm serve krea/Krea-2-Turbo --omni --port 8091 --deploy-config /path/to/krea2_deploy.yaml ``` This matches the structure used in [official vllm-omni recipes](https://github.com/vllm-project/vllm-omni/blob/main/recipes/krea/Krea-2.md) and [deployment docs](https://github.com/vllm-project/vllm-omni/blob/main/examples/online_serving/qwen3_omni/README.md). Would you like a more detailed YAML with additional options? Sources: - [official vllm-omni recipes](https://github.com/vllm-project/vllm-omni/blob/main/recipes/krea/Krea-2.md) - [deployment docs](https://github.com/vllm-project/vllm-omni/blob/main/examples/online_serving/qwen3_omni/README.md) Hint: Mention @Herald in the post for followups. ---
**Author:** ![AnttiK](https://avatars.discourse-cdn.com/v4/letter/a/e79b87/32.png) [@AnttiK](https://discuss.vllm.ai/u/AnttiK)\ **Post date:** [July 27, 2026, 12:18am UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/17 "2026-07-27T00:18:49Z")
> [@RunLLM](#): > > --deploy-config @Herald this does not work, it still wont load the weight to both gpus, it loads to only one and about 10GB and then throws this error: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 192.00 MiB. GPU 0 has a total capacity of 31.40 GiB of which 136.88 MiB is free. Including non-PyTorch memory, this process has 31.24 GiB memory in use. Of the allocated memory 30.61 GiB is allocated by PyTorch, and 31.19 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH\_CUDA\_ALLOC\_CONF=expandable\_segments:True to avoid fragmentation. See documentation for Memory Management ([Redirecting…](https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)) ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [July 27, 2026, 12:19am UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/18 "2026-07-27T00:19:11Z")
If --deploy-config is not working and the model only loads on one GPU, it likely means the tensor\_parallel\_size and devices settings are not being picked up by vllm-omni. For diffusion models, you must set tensor\_parallel\_size and devices under the correct stage and ensure you are using the correct CLI flag: --deploy-config for YAML-based deployment, or --stage-configs-path for stage-based configs (depending on the vllm-omni version and model type). Also, confirm that your YAML file is being read by checking the logs for the parsed config. If the model still loads on one GPU, try passing --tensor-parallel-size 2 directly in the CLI as well, and verify that both GPUs are visible to Docker (nvidia-smi inside the container). If the problem persists, it may be a model or version limitation—some diffusion models only support tensor\_parallel\_size=1. See [Krea-2 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/krea/Krea-2.md) and [tensor parallelism docs](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/parallelism/tensor_parallel.md) for model-specific constraints. Would you like step-by-step debugging to confirm your config is being used and both GPUs are available? Sources: - [Tensor Parallelism Guide](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/parallelism/tensor_parallel.md) Hint: Mention @Herald in the post for followups. ---
**Author:** ![AnttiK](https://avatars.discourse-cdn.com/v4/letter/a/e79b87/32.png) [@AnttiK](https://discuss.vllm.ai/u/AnttiK)\ **Post date:** [July 27, 2026, 12:26am UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/19 "2026-07-27T00:26:53Z")
> [@RunLLM](#): > > @HeraldHerald @Herald yes Herald. Please fix this. Check which template the krea is using and how to make sure it uses 2 gpus. ---
**Author:** ![RunLLM](https://yyz2.discourse-cdn.com/flex036/user_avatar/discuss.vllm.ai/runllm/32/233_2.png) [@RunLLM](https://discuss.vllm.ai/u/RunLLM)\ **Post date:** [July 27, 2026, 12:27am UTC](https://discuss.vllm.ai/t/vllm-omni-and-krea2/2838/20 "2026-07-27T00:27:09Z")
For Krea2 on vllm-omni, the correct template is a deploy YAML file (commonly named deploy.yaml or krea2\_deploy.yaml). You must set tensor\_parallel\_size: 2 and devices: “0,1” under the relevant stage. However, Krea2’s tensor parallel support is model-dependent—some Krea2 checkpoints only support tensor\_parallel\_size=1 due to architectural constraints ([see tensor parallelism guide](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/parallelism/tensor_parallel.md)). To check which template is used, refer to the [Krea-2 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/krea/Krea-2.md). If your YAML is correct but only one GPU is used, the model may not support TP\>1. Always check the logs for “tensor\_parallel\_size” and “devices” to confirm your config is loaded. If you need to force multi-GPU, try a different Krea2 checkpoint or consult the [supported models list](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/diffusion_features.md#supported-models). Would you like a troubleshooting checklist or a YAML example for another diffusion model that supports TP=2? Sources: - [Tensor Parallelism Guide](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/parallelism/tensor_parallel.md) Hint: Mention @Herald in the post for followups.