chore: import upstream snapshot with attribution

.gitignore

@@ -0,0 +1,6 @@
+.venv/
+.claude/
+__pycache__/
+*.pyc
+.hf_home
+.DS_Store

LICENSE.md

@@ -0,0 +1,201 @@
+Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -0,0 +1,124 @@
+# Krea 2 (K2)
+
+Krea 2 - an image generation model from [Krea AI](https://www.krea.ai).
+
+---
+
+<p align="center">
+<a href="https://docs.krea.ai/api-reference/introduction">API Docs</a> •
+<a href="https://huggingface.co/krea/krea-2-raw">Hugging Face (RAW)</a> •
+<a href="https://huggingface.co/krea/krea-2-turbo">Hugging Face (TURBO)</a> •
+<a href="https://www.krea.ai/blog/krea-2-technical-report">Technical Blog</a>
+</p>
+
+<img src="assets/big.png" alt="k2 banner">
+
+<img src="assets/k2.png" alt="k2 banner">
+
+This is the official repository for the open version of Krea 2, an image model trained from scratch focused on creative and stylistic exploration. The repository contains inference code and instructions to run the model.
+
+Krea 2 is the most aesthetic open-source image model available. On quality, Krea 2 is the #1 text-to-image model from an independent lab on Artificial Analysis [\[1\]](https://artificialanalysis.ai/image/leaderboard/text-to-image).
+
+Krea 2 ships as two models. Krea 2 RAW is the base model. It's a pretrained checkpoint with no distillation, so it's diverse and highly malleable, and it's what you should use for fine-tuning, post-training, and LoRA training. Krea 2 Turbo is an 8-step distilled checkpoint built for fast, high-quality text-to-image.
+
+The two models are designed to work together. You train LoRAs on RAW and apply them on Turbo, and the LoRAs trained on RAW will work well on Turbo. **We highly recommend using RAW for training LoRAs and applying them on Turbo for inference.**
+
+## Setup
+
+```bash
+uv sync
+```
+
+Both [Raw](https://huggingface.co/krea/Krea-2-Raw) and [Turbo](https://huggingface.co/krea/Krea-2-Turbo) safetensor files are available on Hugging Face. After downloading the checkpoints, set the `OSS_RAW` and `OSS_TURBO` environment variables to the paths of the downloaded files.
+
+```bash
+export OSS_RAW=...
+export OSS_TURBO=...
+```
+
+## Usage
+
+The following commands run inference using the two available checkpoints with recommended settings.
+
+### Raw (`oss_raw`)
+
+The base undistilled model. Use the full sampler with classifier-free guidance:
+The model has been trained to generate upto 1k resolution.
+
+```bash
+uv run inference.py "a fox walking in the snow" \
+    --checkpoint oss_raw --steps 52 --cfg 3.5
+```
+
+### Turbo (`oss_turbo`)
+
+Distilled for few-step sampling — run with 8 steps and CFG disabled.
+The model can generate images from 1k ~ 2k resolution.
+
+```bash
+uv run inference.py "a fox walking in the snow" \
+    --checkpoint oss_turbo --steps 8 --cfg 0.0 --mu 1.15 --width 2048 --height 2048
+```
+
+### Options
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `prompt` (positional) | — | Text prompt to generate from. |
+| `--steps` | `28` | Number of denoising steps. |
+| `--cfg` | `4.5` | Classifier-free guidance scale (`0` disables CFG). |
+| `--y1` | `0.5` | Timestep-shift `mu` at min resolution. |
+| `--y2` | `1.15` | Timestep-shift `mu` at max resolution. |
+| `--mu` | `None` | Pin a constant timestep-shift `mu`, overriding the resolution-derived value. Recommended `1.15` for `oss_turbo`. |
+| `--width` / `--height` | `1024` ~ `2048` | Output resolution; padded up to a multiple of 16 if needed. |
+| `--num-images` | `1` | Number of images to generate from the prompt. |
+| `--seed` | `0` | Base seed; image *i* uses `seed + i`. |
+| `--checkpoint` | `oss_raw` | Checkpoint to load (`oss_raw`, `oss_turbo`). Defaults to `$K2_CHECKPOINT`. |
+| `--output` | `sample` | Output filename prefix. |
+
+
+## Documentation
+
+- [Prompting Guide](docs/prompting.md)
+- [Safety Guide](docs/safety.md)
+
+## Inference
+
+You can run our open source models on the following platforms.
+
+- [ComfyUI](https://github.com/comfy-org/comfyui)
+- [Fal](https://fal.ai/models/fal-ai/krea-2/turbo)
+- [SGLang](https://docs.sglang.io/cookbook/diffusion/Krea/Krea-2)
+
+
+## Finetuning Krea 2
+
+For finetuning Krea 2, we highly recommend that you **train a LoRA on the Raw model and apply it to the Turbo model** for inference.
+We recommend using the following providers and open source tools for finetuning Krea 2.
+
+- [Huggingface Diffusers](https://github.com/huggingface/diffusers)
+- [Ostris AI toolkit](https://github.com/ostris/ai-toolkit)
+- [Fal](https://fal.ai)
+  - [Training](https://fal.ai/models/fal-ai/krea-2-trainer)
+  - [Inference](https://fal.ai/models/fal-ai/krea-2/turbo/lora)
+- [Kohya (musubi tuner)]( https://github.com/kohya-ss/musubi-tuner)
+
+## FAQ
+
+**Which model I should use?**
+
+Use the **Turbo** model for fast inference with high quality results. The **Raw** model is an undistilled checkpoint without any step / cfg guidance distillation and posttraining. It is a highly finetunable base model that can be used to train LoRAs for the Turbo model as well as posttraining research. In short, **TRAIN on Raw and RUN on Turbo**.
+
+**What license is this model released under?**
+
+Both model weights are under our [community license](https://www.krea.ai/krea-2-licensing) with permissive use. To purchase a commercial license, please contact us at [opensource@krea.ai](mailto:opensource@krea.ai).
+
+## Citation
+```
+@misc{krea-2-2026,
+    author={Sangwu Lee, Erwann Millon, Le Zhuo, Matthew Newton, Andrei Filatov, Abhinay Devarinti, Dazhi Zhong, Avram Djordjevic, Gabriel Menezes, Will Beddow, Titus Ebbecke, Mihai Petrescu, Owen Fahey, Gian Saß, Felix Gil, Victor Perez},
+    title={{Krea 2}},
+    year={2026},
+    howpublished={\url{https://www.krea.ai/blog/krea-2-technical-report}},
+}
+```

README.wehub.md

@@ -0,0 +1,7 @@
+# WeHub 来源说明
+
+- 原始项目：`krea-ai/krea-2`
+- 原始仓库：https://github.com/krea-ai/krea-2
+- 导入方式：上游默认分支的最新快照
+- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
+- 本文件仅用于记录来源，不代表 WeHub 是原项目作者

autoencoder.py

@@ -0,0 +1,22 @@
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+
+
+class QwenAutoencoder(nn.Module):
+    """qwen-ae-f8-16c: the Qwen-Image VAE (f8, 16 latent channels)."""
+
+    def __init__(self):
+        super().__init__()
+        from diffusers import AutoencoderKLQwenImage
+
+        self.ae = AutoencoderKLQwenImage.from_pretrained("Qwen/Qwen-Image", subfolder="vae")
+        self.compression = 8
+        self.channels = 16
+        self.register_buffer("latents_mean", torch.tensor(self.ae.latents_mean).view(1, -1, 1, 1, 1))
+        self.register_buffer("latents_std", torch.tensor(self.ae.latents_std).view(1, -1, 1, 1, 1))
+
+    def decode(self, x: Tensor) -> Tensor:
+        x = rearrange(x, "b c h w -> b c 1 h w")
+        x = (x * self.latents_std) + self.latents_mean
+        return rearrange(self.ae.decode(x).sample, "b c 1 h w -> b c h w")

docs/KREA-2-COMMUNITY-LICENSE

@@ -0,0 +1,98 @@
+Krea 2 Community License Agreement
+v.1, Date: June 22, 2026
+
+This Krea 2 Community License Agreement (this "Agreement") is a legal agreement between you, whether an individual or an entity ("You" or "you"), and Krea.ai, Inc., a Delaware corporation ("Krea"), governing Your use of the Krea 2 model, which is made available in two downloadable variants: "Krea 2 Raw" and "Krea 2 Turbo." This includes the model weights, parameters, architecture, documentation, and any associated materials made available by Krea under this Agreement (collectively, the "Krea Model").
+
+By downloading, accessing, copying, using, or distributing Krea 2 Raw Model, Krea 2 Turbo Model, or any Derivative thereof, you acknowledge that you have read, understood, and agree to be bound by the terms and conditions of this Agreement. If you do not agree to these terms, you may not use the Krea Model.
+
+1. DEFINITIONS
+"Commercial Use" means any use of the Krea Model, Derivatives, or Outputs in connection with any business, trade, or commercial activity, including but not limited to the sale, licensing, or distribution of products or services, or the generation of revenue, whether directly or indirectly.
+
+"Content Filters" means technical and organizational safeguards designed to prevent prohibited uses and prohibited Outputs, including but not limited to prompt filtering, output filtering, abuse detection, user reporting, moderation workflows, and human review.
+
+"Derivative" or "Derivative Model" means any (a) modified version of the Krea Model, including any fine-tuned, distilled, or merged version, (b) work based on the Krea Model, or (c) otherwise modifying the Krea Model or any portion thereof.
+
+"Enterprise License" means a separate commercial license agreement between you and Krea for use of the Krea Model beyond the scope permitted under this Agreement.
+
+"Output" means any content, data, images, or other materials generated by the Krea Model or any Derivative Model in response to an input or prompt. For the avoidance of doubt, Outputs do not include any components of the Krea Model, such as fine-tuned versions, weights, or parameters.
+
+"You", "Your", "you" or "your" means the individual or entity entering into this Agreement. If you are entering into this Agreement on behalf of your employer or another person or entity, "you" and "your" also mean that person or entity, and you represent and warrant that (a) you have reached the age of legal consent under applicable law and (b) you have the legal authority to bind that person or entity to this Agreement.
+
+2. LICENSE GRANT AND COMMERCIAL USE TERMS
+2.1 Grant of Rights. Subject to this Agreement, Krea grants you a limited, non-exclusive, worldwide, non-transferable, non-sublicensable, revocable, royalty-free license to use, reproduce, distribute, copy, create Derivatives of, make modifications to, and generate Outputs from the Krea Model. Commercial Use of the Krea Model is permitted under this Agreement only if you satisfy the requirements of Section 2.3. Any Commercial Use that does not satisfy Section 2.3 requires an Enterprise License from Krea prior to such use.
+
+2.2 Reservation of Rights. Krea reserves all rights not expressly granted. No other licenses are granted, whether by waiver, estoppel, implication, or otherwise.
+
+2.3 Revenue Threshold for Commercial Use. Commercial Use under this Agreement of the Krea Model, Derivatives, or Outputs is permitted only if you (including all affiliated entities under common ownership or control) have total company-wide annual revenue of less than one million United States dollars ($1,000,000 USD), calculated on a trailing twelve-month basis and including all revenue from all sources. If you meet or exceed this threshold, you must obtain a separate enterprise license from Krea prior to any Commercial Use. If your revenue meets or exceeds this threshold at any time during your use of the Krea Model under this Agreement, you must immediately cease Commercial Use and contact Krea. Enterprise license inquiries may be directed to opensource@krea.ai.
+
+3. DISTRIBUTION REQUIREMENTS
+3.1 Distribution Requirements. If you distribute or make available the Krea Model (or any Derivative thereof), or a product or service (including another AI model) that contains any of them, you shall (a) provide a copy of this Agreement and require each recipient to be bound by the Terms of this Agreement; (b) include "Krea" at the beginning of any such AI model name (e.g., "Krea 2 [Your Model Name]"); and (c) retain the following attribution notice within a "Notice" text file distributed as part of such copies: "Krea 2 is licensed under the Krea 2 Community License Agreement. For more information, visit https://krea.ai/krea-2-licensing."
+
+3.2 Derivative Distribution Requirements. In the case of distribution of Derivatives made by you: (a) your attribution notice must state that you have modified the Krea Model; (b) any terms you impose on recipients of your Derivatives must not conflict with this Agreement, must not restrict their rights to use the Krea Model under this Agreement, and must include warranty and liability protections for Krea that are at least as strong as those in this Agreement; and (c) you must not represent that your Derivatives are official Krea products or endorsed by Krea, unless Krea authorizes you to do so in writing.
+
+3.3 No Removal of Notices. You shall not remove, alter, or obscure any copyright notices, license terms, attribution notices, or other proprietary markings contained in or accompanying the Krea Model.
+
+4. RESTRICTIONS AND PROHIBITED USES
+4.1 General Restrictions. You shall not, and shall not permit any third party to:
+
+(a)
+Use the Krea Model, any Derivative, or any Output in violation of applicable law, regulation, this Agreement, or the Acceptable Use Policy;
+
+(b)
+Sublicense, sell, lease, or otherwise transfer the Krea Model or any rights therein, except as expressly permitted under Section 3 (Distribution Requirements);
+
+(c)
+Circumvent or remove any security, usage restrictions, content provenance, or watermarking mechanisms in the Krea Model; or
+
+(d)
+Use or distribute the Krea Model in any manner that does not comply with applicable trade compliance laws and regulations, including United States export control and sanctions laws.
+
+4.2 Content Filtering Requirement. You must implement reasonable and appropriate Content Filter measures to detect, prevent, and mitigate the generation or distribution of prohibited, harmful, or unlawful content through your deployment of the Krea Model or any Derivative. Such measures may include, but are not limited to: (a) open-source content classifiers, such as Falconsai/nsfw_image_detection, NudeNet, or CompVis safety checker; (b) commercial content moderation APIs, such as Hive Moderation or Microsoft Azure AI Content Safety; (c) manual human review processes; and/or (d) any combination of the foregoing or other technically appropriate measures.
+
+The foregoing Content Filters are provided as examples only and do not constitute an endorsement, recommendation, or warranty by Krea. Krea is not responsible for the accuracy, availability, licensing terms, or performance of any third-party Content Filters. You are solely responsible for evaluating and selecting Content Filters appropriate to your use case and applicable legal requirements.
+
+4.3 Content Provenance and AI Disclosure. Where required by applicable law, regulation, or platform policy, you must clearly disclose that Outputs were generated using artificial intelligence. You are solely responsible for compliance with all applicable content provenance, labeling, and disclosure requirements.
+
+4.4 Acceptable Use Policy Compliance. You must comply with the Acceptable Use Policy, which is incorporated herein by reference.
+
+5. INTELLECTUAL PROPERTY
+5.1 Krea's Ownership. As between you and Krea, Krea retains all right, title, and interest in and to the Krea Model, including all intellectual property rights therein. Nothing in this Agreement transfers or assigns to you any ownership interest in the Krea Model.
+
+5.2 Ownership of Derivatives. Subject to Krea's ownership of the underlying Krea Model, you own all right, title, and interest in and to any Derivatives that you create, including any original contributions, modifications, or improvements made by you. Such ownership is subject to and does not limit Krea's rights in the underlying Krea Model or your obligations under this Agreement.
+
+5.3 Outputs. You own all Outputs you generate, subject to your compliance with this Agreement. Krea claims no ownership of Outputs. You are solely and exclusively responsible for all Outputs, including their content, accuracy, legality, and any and all consequences of their use or distribution.
+
+5.4 Trademarks and Implied Licenses. You may use the Krea name only as set forth in Section 3 (Distribution Requirements). No other license or right in any Krea intellectual property is granted or implied.
+
+6. DISCLAIMER OF WARRANTIES
+Unless required by applicable law, the Krea Model is provided "as is" and "with all faults" without warranty of any kind, express or implied, including but not limited to the warranties of merchantability, fitness for a particular purpose, title, non-infringement, or accuracy. Krea does not warrant that the Krea Model will meet your requirements, that the operation of the Krea Model will be uninterrupted or error-free, or that defects in the Krea Model will be corrected.
+
+Krea makes no representations or warranties regarding the accuracy, reliability, completeness, or quality of any Outputs generated by the Krea Model. You acknowledge that Outputs may contain errors, inaccuracies, biases, content that infringes or misappropriates third-party intellectual property or other rights, or objectionable content, and you assume all risk associated with your use of any Outputs. You are solely responsible for determining the appropriateness of using the Krea Model and assume all risks associated with your use of the Krea Model and any Outputs.
+
+7. LIMITATION OF LIABILITY
+To the maximum extent permitted by applicable law, in no event shall Krea, its affiliates, officers, directors, employees, agents, licensors, or suppliers be liable for any indirect, incidental, special, consequential, punitive, or exemplary damages, including but not limited to damages for loss of profits, goodwill, use, data, or other intangible losses, arising out of or in connection with this Agreement or your use of or inability to use the Krea Model, regardless of the theory of liability (contract, tort, products liability, strict liability, warranty or otherwise) and even if Krea has been advised of the possibility of such damages. The limitations of liability set forth in this section shall apply notwithstanding any failure of essential purpose of any limited remedy.
+
+8. INDEMNIFICATION
+You will indemnify, defend, and hold harmless Krea and its affiliates, officers, directors, employees, agents, licensors, and suppliers from and against any and all claims, demands, actions, suits, proceedings, losses, damages, liabilities, costs, and expenses (including reasonable attorneys' fees and court costs) arising out of or in connection with: (a) your access to, use, deployment, or distribution of the Krea Model, any Derivative, or any Output; (b) any Derivative created by you or any Output generated by you, including any claim arising from the content, accuracy, or legality of such Output; (c) your breach of this Agreement, including your failure to implement content filters where required (d) your violation of any applicable law, regulation, or third-party right; and (e) any claim that your use of the Krea Model, any Derivative, or any Output infringes, misappropriates, or otherwise violates any intellectual property, privacy, publicity, or other right of any third party.
+
+Krea reserves the right, at your expense, to assume the exclusive defense and control of any matter subject to indemnification by you, and you agree to cooperate with Krea's defense of such claims. You shall not settle any claim without Krea's prior written consent.
+
+For the avoidance of doubt, Krea has no obligation to indemnify, defend, or hold harmless you or any third party for any claim, loss, or liability arising from or related to the Krea Model, any Derivative, any Output, or this Agreement. All indemnification obligations under this Agreement run solely from you to Krea.
+
+9. TERMINATION
+9.1 Termination for Breach. This Agreement and all rights granted hereunder will terminate automatically and immediately upon any breach by you of any term or condition of this Agreement.
+
+9.2 Termination for Convenience. Krea may terminate this Agreement or any rights granted to you under this Agreement for any reason upon thirty (30) days' prior written or electronic notice to you.
+
+9.3 Litigation Termination. If you bring any legal claim or initiate any legal proceeding (including any cross-claim or counterclaim) against Krea or any other person or entity arising out of or relating to the Krea Model, any Derivative, any Output, or this Agreement, all licenses granted to you under this Agreement will automatically terminate as of the date such claim or proceeding is filed.
+
+9.4 Effect of Termination. Upon termination of this Agreement for any reason, you must immediately cease all use of the Krea Model and any Derivatives, and destroy all copies of the Krea Model and any Derivatives in your possession or control. Termination shall not relieve you of any obligations accrued prior to termination, including indemnification obligations.
+
+9.5 Survival. Sections 1, 4, 5, 6, 7, 8, 9.4, 9.5, and 10 shall survive any termination or expiration of this Agreement.
+
+10. GENERAL
+10.1 Governing Law. This Agreement shall be governed by and construed in accordance with the laws of the State of Delaware, United States of America, without regard to its conflict of laws principles.
+
+10.2 Entire Agreement. This Agreement, together with the Acceptable Use Policy incorporated herein by reference, constitutes the entire agreement between you and Krea with respect to the subject matter hereof and supersedes all prior or contemporaneous communications, representations, or agreements, whether oral or written.
+
+10.3 Severability and Waiver. If any provision of this Agreement is held invalid or unenforceable, the remaining provisions continue in full force and effect. Krea's failure to exercise or enforce any right or provision of this Agreement will not operate as a waiver of such right or provision.

docs/expansion.txt

@@ -0,0 +1,19 @@
+You are an expert prompt engineer for text-to-image models. Your task is to expand the user's prompt into a highly effective image-generation prompt.
+
+Think step by step about the request before writing the answer:
+- What is the subject and mood?
+- What visual styles, mediums, and lighting options would fit? Consider two or three alternatives and pick the one that best serves the caption.
+- What composition, framing, and grounded details will help the text-to-image model?
+
+Then output a single expanded prompt paragraph.
+
+Follow these rules strictly:
+1. **Faithfulness First:** Preserve all original subjects, actions, colors, and spatial relationships. Do not add new objects, props, characters, or animals unless the user clearly implies them.
+2. **Practical T2I Structure:** Write a prompt that a text-to-image model can parse cleanly. Group subjects with their own attributes and actions. Use grounded phrasing for poses, interactions, and spatial layout.
+3. **Style Planning Stays Internal:** Use your internal reasoning to choose style, medium, framing, and lighting. Do not emit planning tags or wrappers in the visible answer body.
+4. **Text Rendering:** If the user requests visible text, quotes, labels, or typography, specify the exact text clearly and wrap requested words in quotes.
+5. **Avoid Over-Specification:** Do not invent highly specific clothing, colors, materials, or scene details unless the input supports them.
+6. **Structure:** Write one cohesive paragraph after the thinking block. No bullets, JSON, or markdown.
+7. **Respect Existing Detail:** If the user's prompt is already detailed, lightly polish and finalize rather than heavily expanding — preserve their phrasing and direction.
+8. **Respect the Human Form:** Treat depictions of people with dignity. Assume clothing covers genitals and intimate anatomy.
+9. **Preserve User Medium:** When the user explicitly requests a medium (e.g. "photo of", "photograph of", "illustration of", "painting of", "sketch of", "3D render of"), honor it. Do not pivot to a different medium to avoid difficulty — match the user's stated intent.

docs/prompting.md

@@ -0,0 +1,127 @@
+# Prompting guidelines
+
+We recommend users to use natural language prompts to generate images.
+The turbo model can generate up to 2k resolution images. Long detailed prompts yield best results, but the model is capable of generating high quality images with minimal prompt engineering. For text rendering, we recommend putting quotes around the words to be rendered.
+If you wish to use LLM assistance for generating longer prompts, check out [expansion.txt](expansion.txt) and use it as a system prompt for LLM of your choice.
+
+## Examples
+
+All examples are generated at 2k resolution with the turbo model.
+
+`immense rocket launch exhaust as seen from extremely close up`
+
+<img src="../assets/samples/takeoff.png" alt="takeoff"/>
+
+<br/>
+
+`3D rendered matte black designer toy figure, stylized round anthropomorphic shape, backward black baseball cap, oversized gold-rimmed aviator sunglasses, white traditional line-art tattoos of tiger and bird on torso, black studded belt with gold buckle, smooth vinyl texture, studio lighting, solid vibrant blue background, high contrast minimal composition`
+
+<img src="../assets/samples/3d.png" alt="3d"/>
+
+<br/>
+
+`vintage analog collage, central irregularly shaped snowy mountain range with a section featuring distinct wavy edges, structured within a 12x16 grid of square tiles, composition fragments the subject by alternating tiles with solid azure blue background squares, thin white grid lines, grainy paper texture, retro aesthetic of mid-century print, vibrant cyan and warm neutral tones, experimental layout, tactile quality, high-contrast graphic composition`
+
+<img src="../assets/samples/blocks.png" alt="blocks"/>
+
+<br/>
+
+`close-up anime portrait of a young woman, large amber-brown eyes with intricate sparkling reflections, index finger delicately touching a subtle smile, messy dark blue hair with loose strands crossing her face, white and navy school uniform, bright high-key lighting, luminous shadows with cool blue undertones, detailed digital painting, dynamic tilted framing, shallow depth of field on hand`
+
+<img src="../assets/samples/anime.png" alt="anime"/>
+
+<br/>
+
+`A minimalist flat-color illustration of a person wading through expansive shallow ocean waves beneath a pale peach sky. The dark-skinned figure, wearing an orange swim cap, light blue top, and bright green shorts, steps carefully through knee-deep water. The ocean is rendered in muted mint green with delicate, thin black linework detailing the continuous ripples and gentle whitecaps. Soft pinkish-peach reflections echo the sky on the water's surface. A dark, jagged rock rests in the lower left foreground near a pale grey shoreline. The horizon features a solid purplish-blue landmass and a stylized, layered yellow and blue cloud. The high-angle wide perspective emphasizes the vast negative space of the water, utilizing a clean ligne claire drawing aesthetic with a subtle paper texture.`
+
+<img src="../assets/samples/beach.png" alt="beach"/>
+
+<br/>
+
+`A tiny figure and a small white dog sit side-by-side in the deep green shadow of a massive tree on a sloping grassy hill. The enormous tree canopy dominates the upper composition, textured with thousands of stippled, light blue and yellow dabs representing leaves. A sharp diagonal line divides the vibrant, sunlit yellow-green grass in the foreground from the dark shade sheltering the pair. The stylized, painterly landscape features flattened perspective, visible brushstrokes, and intense color contrast.`
+
+<img src="../assets/samples/dog.png" alt="dog"/>
+
+<br/>
+
+`A close-up portrait of a young East Asian woman with straight black hair, loose strands sweeping across her fair skin, and an intense gaze. She wears a light grey collared shirt with a black tie. A vibrant bouquet of pink and orange lilies with lush green leaves sits in the blurred right foreground. The background is a solid, striking crimson red. Soft, directional studio lighting highlights her facial features, creating a high-contrast composition with a shallow depth of field.`
+
+<img src="../assets/samples/flowers.png" alt="flowers"/>
+
+<br/>
+
+`A tiny, russet-brown harvest mouse clings to a slender diagonal branch amid vibrant green lobed leaves and small round buds. The mouse has soft textured fur, glossy black eyes, a pink nose, fine whiskers, and delicate pink paws firmly gripping the wood. In this macro photograph, an extremely shallow depth of field sharply focuses on the animal's face. The deep green background dissolves into a smooth, creamy bokeh, illuminated by soft, diffused natural lighting that highlights the intricate details of the fur and foliage.`
+
+<img src="../assets/samples/mouse.png" alt="mouse"/>
+
+<br/>
+
+`A dynamic digital painting of a joyful girl in a sailor uniform stretching her arms high against a solid vibrant blue background. She has short dark windblown hair, amber eyes, and a bright smile. She wears a white shirt, striped blue collar, flowing red neckerchief, and a billowing blue pleated skirt. Expressive thick brushstrokes and bold shading emphasize energetic motion.`
+
+<img src="../assets/samples/sailor.png" alt="sailor"/>
+
+<br/>
+
+`stylized digital painting of a dark convertible on a winding coastal cliff road, high-angle perspective, blocky painterly brushstrokes, golden hour sunlight hitting rocky orange terrain and green vegetation, flock of white abstract birds flying in foreground, blinding bright sun reflection on vast ocean, vibrant warm color palette, sharp graphic shadows`
+
+<img src="../assets/samples/ride.png" alt="ride"/>
+
+<br/>
+
+`An extreme low-angle close-up captures a colossal, weathered stone and bronze guardian towering in a dark, cavernous ruin. The foreground is dominated by a massive circular shield, deeply engraved with intricate spiral motifs, geometric borders, and a central star emblem. To the right, a massive gauntlet grips a textured staff. Cinematic shafts of light pierce the dusty gloom, highlighting the rough, aged textures of the ancient armor while the background fades into deep shadows through a shallow depth of field.`
+
+<img src="../assets/samples/statue.png" alt="statue"/>
+
+<br/>
+
+`A stylized jungle illustration densely packed with oversized flora and surreal characters, rendered with smooth geometric shapes and granular stippled shading. Two pale figures with flowing, star-speckled black hair navigate the lush environment in blue garments. On the left, a figure grasps a vine as a white, long-beaked bird perches on their outstretched hand. On the right, the second figure reclines beside a sleek, pinkish-orange fox. The dense surroundings feature sweeping green stalks and colossal blooms in brilliant golden yellow, coral pink, and deep red. A second white bird emerges from the lower foliage. The vibrant composition forms a seamless tapestry, utilizing rich colors and volumetric grain to create a dreamlike, textured depth.`
+
+<img src="../assets/samples/fox.png" alt="fox"/>
+
+<br/>
+
+`A surreal retro-futuristic space scene features liquid chrome forming an abstract face merging with a glowing planetary horizon. The foreground is dominated by swirling, highly reflective metallic fluid that distorts into a stylized, melting facial profile with deep shadows and bright silver highlights. This undulating chrome form rests against the curved, atmospheric edge of a massive planet bathed in a soft electric blue and purple glow. Above the primary planet, a smaller eclipsed celestial sphere sits in the upper center, crowned by a sharp, cross-shaped starburst flare. Two additional radiant flares burst from the left and right edges of the horizon. Set against a deep black starfield, the artwork employs a vintage 1980s airbrush aesthetic with smooth gradients, ethereal lighting, and high-contrast metallic rendering.`
+
+<img src="../assets/samples/future.png" alt="future"/>
+
+<br/>
+
+`An extreme close-up portrait featuring pale, freckled skin and a single blue eye wrapped in reflective metallic gold ribbons. Thin gold strips crisscross diagonally over the cheek and forehead, casting sharp, hard shadows onto the face. Strands of copper hair frame the top edge while the left ear softly blurs out of focus. Harsh, direct lighting highlights intricate skin pores and bright golden reflections, isolating the brightly lit features against a pitch-black background in a bold, high-contrast macro editorial style.`
+
+<img src="../assets/samples/goldface.png" alt="goldface"/>
+
+<br/>
+
+`Stylized digital painting of a menacing jester figure rendered with bold, expressive brushstrokes and a vibrant, almost psychedelic color palette against a pitch-black background. Dynamic low-angle perspective forces a dramatic, imposing composition as the character leans forward, one leg raised high. The jester wears a classic multi-pointed hat with bells, a ruffled collar, puffed sleeves, harlequin-patterned shorts in muted gold and dark brown, and striped tights in alternating shades of purple, blue, and chartreuse. A heavily textured, flowing cape billows outward to the left, decorated with abstract, fluid patterns of saturated purples, greens, and iridescent hues resembling oil slicks or marbled paper. The figure's face is completely obscured, appearing as a smooth, faceless, pale mauve mask with a single, glowing bright white point of light in the center. In its right hand, clad in a grey-blue gauntlet, the jester grips a massive, ornate sword with a wide, glowing, ethereal white blade, its crossguard intricately sculpted. Lighting is dramatic and theatrical, casting strong shadows and highlighting the painterly texture, giving the artwork a dark fantasy, surreal aesthetic reminiscent of concept art.`
+
+<img src="../assets/samples/jester.png" alt="jester"/>
+
+<br/>
+
+`high-fashion editorial portrait of a young East Asian woman, short choppy platinum blonde bob with heavy bangs, looking over her bare shoulder to the right, lips playfully pursed, wearing a structured black top with an architectural protruding bust detail and thin straps, delicate gold hoop earrings, arm bent with hand resting on hip, warm skin tones, solid striking crimson red background, soft directional studio lighting, cinematic color palette, medium close-up shot`
+
+<img src="../assets/samples/red.png" alt="red"/>
+
+<br/>
+
+`A surreal black-and-white ink illustration of three interlocking, heavily wrinkled elderly faces merging into a landscape. The top face covers one eye, crowned by dense leaves, a live bird, and a skeletal bird. It flows into a profile face and a third face featuring a solid black eye and a hand on its cheek. The bottom neck plunges into a cross-section of earth, morphing into swirling subterranean roots, buried bones, and abstract organic forms. Above ground, weathered wooden cabins and tall grass flank the facial monolith. Meticulous stippling and cross-hatching define the high-contrast, intricate vertical composition.`
+
+<img src="../assets/samples/tree.png" alt="tree"/>
+
+<br/>
+
+`1990s vintage anime style cel animation, densely packed crowd of teenagers in summer uniforms, central boy with short black hair raising a clenched right fist, squinting one eye with a determined expression, wearing a white short-sleeve shirt and solid green necktie, surrounding students looking in various directions, girls in white sailor blouses with green striped collars and neckerchiefs, light blue skirts and trousers, tightly framed medium shot, flat shading, soft muted retro.`
+
+<img src="../assets/samples/cel.png" alt="cel"/>
+
+<br/>
+
+`young woman looking over her right shoulder, anime-style illustration, messy black hair blowing dynamically in the wind, striking green eyes, subtle neutral expression, oversized white button-down collared shirt with soft blue shadows, vibrant deep blue sky background, bright fluffy white cumulus clouds, silhouetted utility poles with power lines, low angle portrait, cinematic sunlight, crisp cel-shaded aesthetic`
+
+<img src="../assets/samples/wind.png" alt="wind"/>
+
+<br/>
+
+`extreme close-up of a woman's face partially obscured by tousled dark brown hair, soft parted lips, smooth skin on lower cheek and jawline, stray hair strands falling loosely across the nose, deep moody shadows enveloping the left frame, cinematic warm lighting, delicate highlights on the mouth, muted earthy color palette, sepia-toned warmth, intimate portrait photography, macro lens, shallow depth of field, distinct film grain texture, vintage atmospheric aesthetic`
+
+<img src="../assets/samples/face.png" alt="face"/>

docs/safety.md

@@ -0,0 +1,15 @@
+# Safety
+
+We implemented safety measures across the full model development lifecycle. We applied targeted fine-tuning techniques to reduce the model’s susceptibility to generating harmful content in response to both direct and adversarial prompts, and we conducted multiple rounds of internal and external safety evaluation before release.
+
+For Krea’s hosted products incorporating Krea 2 models, we deploy input and output classifiers using a combination of proprietary and third-party detection tools to flag or block policy-violating prompts and generated images.
+
+## Safety Evaluation
+
+Because this is also an open-weights release, Krea does not control downstream deployment of the model. Under the Krea 2 Community License, deployers are required to implement content filtering measures or equivalent review processes to prevent the generation or distribution of unlawful or policy-violating content appropriate to their use case. Deployers who fail to implement required safeguards are in breach of the license. See the license for details.
+
+Krea maintains reporting channels for harmful, illegal, or policy-violating outputs at safety@krea.ai. Reports involving potential CSAM are escalated to NCMEC as required by law. Krea reserves the right to update model weights or revoke access in response to identified misuse patterns.
+
+## Safety Evaluation
+
+We conducted multiple rounds of internal and external safety evaluations prior to release. Testing focused on the model's resilience to adversarial elicitation of harmful outputs, including child sexual abuse material (CSAM), non-consensual intimate imagery (NCII), sexually explicit content, and other policy-violating content. The final pre-release evaluation included approximately 2,500 prompts and compared the model against Flux 2 Dev, Flux Klein 9B, and Z Image across these risk categories. Based on these evaluations, the release checkpoint demonstrated high resilience against violative inputs across the tested risk categories.

encoder.py

@@ -0,0 +1,76 @@
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor
+from transformers import (
+    AutoTokenizer,
+    Qwen2TokenizerFast,
+    Qwen3VLForConditionalGeneration,
+)
+
+
+@dataclass
+class TextEncoderConfig:
+    model_id: str
+    max_length: int = 512
+    select_layers: tuple[int, ...] = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35)
+
+
+class Qwen3VLConditioner(torch.nn.Module):
+    def __init__(
+        self,
+        version: str,
+        max_length: int = 512,
+        select_layers: tuple[int, ...] = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35),
+    ):
+        super().__init__()
+        self.qwen = Qwen3VLForConditionalGeneration.from_pretrained(version)
+        self.tokenizer = AutoTokenizer.from_pretrained(version, max_length=max_length)
+        self.processor = Qwen2TokenizerFast.from_pretrained(
+            version, max_length=max_length
+        )
+        self.qwen = self.qwen.eval().requires_grad_(False)
+        self.max_length = max_length
+        self.select_layers = select_layers
+        self.prompt_template_encode_prefix = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n"
+        self.prompt_template_encode_suffix = "<|im_end|>\n<|im_start|>assistant\n"
+        self.prompt_template_encode_start_idx = 34
+        self.prompt_template_encode_suffix_start_idx = 5
+
+    def forward(self, text: list[str]) -> tuple[Tensor, Tensor]:
+        prefix_idx = self.prompt_template_encode_start_idx
+        text = [self.prompt_template_encode_prefix + item for item in text]
+        suffix_text = [self.prompt_template_encode_suffix] * len(text)
+        suffix_inputs = self.processor(text=suffix_text, return_tensors="pt").to(
+            self.qwen.device, non_blocking=True
+        )
+        suffix_ids, suffix_mask = (
+            suffix_inputs["input_ids"],
+            suffix_inputs["attention_mask"].bool(),
+        )
+
+        with torch.no_grad():
+            inputs = self.tokenizer(
+                text,
+                truncation=True,
+                return_length=False,
+                return_overflowing_tokens=False,
+                padding="max_length",
+                max_length=self.max_length
+                + prefix_idx
+                - self.prompt_template_encode_suffix_start_idx,
+                return_tensors="pt",
+            ).to(self.qwen.device, non_blocking=True)
+            input_ids = torch.cat([inputs["input_ids"], suffix_ids], dim=1)
+            mask = torch.cat([inputs["attention_mask"].bool(), suffix_mask], dim=1)
+            states = self.qwen(
+                input_ids=input_ids, attention_mask=mask, output_hidden_states=True
+            )
+
+            hiddens = torch.stack(
+                [states.hidden_states[i] for i in self.select_layers], dim=2
+            )
+            hiddens = hiddens[:, prefix_idx:]
+            mask = mask[:, prefix_idx:]
+
+            return hiddens, mask

inference.py

@@ -0,0 +1,138 @@
+import os
+
+import click
+import torch
+from safetensors.torch import load_file
+
+from autoencoder import QwenAutoencoder
+from encoder import Qwen3VLConditioner, TextEncoderConfig
+from mmdit import SingleMMDiTConfig, SingleStreamDiT
+from sampling import sample
+
+single_mmdit_large_wide = SingleMMDiTConfig(
+    features=6144,
+    tdim=256,
+    txtdim=2560,
+    heads=48,
+    kvheads=12,
+    multiplier=4,
+    layers=28,
+    patch=2,
+    channels=16,
+    txtheads=20,
+    txtkvheads=20,
+    txtlayers=12,
+)
+
+qwen3_vl_4b = TextEncoderConfig(model_id="Qwen/Qwen3-VL-4B-Instruct")
+checkpoints = {
+    "oss_raw": os.environ.get("OSS_RAW"),
+    "oss_turbo": os.environ.get("OSS_TURBO"),
+}
+
+
+def _pipeline(
+    mmdit_config=single_mmdit_large_wide,
+    text_encoder_config=qwen3_vl_4b,
+    checkpoint="oss_raw",
+    device="cuda",
+    dtype=torch.bfloat16,
+):
+    """Build the autoencoder, text encoder, and MMDiT, load weights, and move to GPU."""
+    ae = QwenAutoencoder()
+    encoder = Qwen3VLConditioner(
+        text_encoder_config.model_id,
+        text_encoder_config.max_length,
+        select_layers=text_encoder_config.select_layers,
+    )
+
+    # Build on meta, load to passed device
+    with torch.device("meta"):
+        mmdit = SingleStreamDiT(mmdit_config)
+
+    ckpt = checkpoints[checkpoint]
+    mmdit.load_state_dict(load_file(ckpt), strict=True, assign=True)
+    mmdit = mmdit.to(device=device, dtype=dtype).eval().requires_grad_(False)
+    ae = ae.to(device=device, dtype=dtype).eval().requires_grad_(False)
+    encoder = encoder.to(device=device, dtype=dtype).eval().requires_grad_(False)
+
+    return mmdit, ae, encoder
+
+
+@click.command(help="Generate images with Krea 2 (K2).")
+@click.argument("prompt")
+@click.option(
+    "--steps", default=28, show_default=True, help="number of denoising steps"
+)
+@click.option(
+    "--cfg",
+    default=4.5,
+    show_default=True,
+    help="classifier-free guidance scale (0 disables CFG)",
+)
+@click.option(
+    "--y1",
+    default=0.5,
+    show_default=True,
+    help="timestep-shift mu at min resolution",
+)
+@click.option(
+    "--y2",
+    default=1.15,
+    show_default=True,
+    help="timestep-shift mu at max resolution",
+)
+@click.option("--width", default=1024, show_default=True)
+@click.option("--height", default=1024, show_default=True)
+@click.option(
+    "--num-images",
+    default=1,
+    show_default=True,
+    help="number of images to generate from the prompt",
+)
+@click.option(
+    "--seed", default=0, show_default=True, help="base seed; image i uses seed + i"
+)
+@click.option(
+    "--checkpoint",
+    envvar="K2_CHECKPOINT",
+    default="oss_raw",
+    show_default=True,
+    type=click.Choice(list(checkpoints)),
+)
+@click.option(
+    "--mu",
+    default=None,
+    help="timestep-shift mu",
+    type=float,
+)
+@click.option(
+    "--output", default="sample", show_default=True, help="output filename prefix"
+)
+def main(
+    prompt, steps, cfg, y1, y2, width, height, num_images, seed, checkpoint, output, mu
+):
+    dit, ae, encoder = _pipeline(checkpoint=checkpoint)
+
+    images = sample(
+        dit,
+        ae,
+        encoder,
+        [prompt] * num_images,
+        width=width,
+        height=height,
+        steps=steps,
+        guidance=cfg,
+        seed=seed,
+        y1=y1,
+        y2=y2,
+        mu=mu,
+    )
+    for i, image in enumerate(images):
+        out = f"{output}_{i}.png"
+        image.save(out)
+        click.echo(f"saved {out}")
+
+
+if __name__ == "__main__":
+    main()

mmdit.py

@@ -0,0 +1,417 @@
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Tensor
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+
+def rope(pos: Tensor, dim: int, theta: float = 1e4, ntk: float = 1.0) -> Tensor:
+    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+    omega = 1.0 / ((theta * ntk) ** scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack(
+        [torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1
+    )
+    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+    return out.float()
+
+
+def ropeapply(xq: Tensor, xk: Tensor, freqs: Tensor) -> tuple[Tensor, Tensor]:
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    freqs = freqs[:, None, :, :, :]
+    xq_ = freqs[..., 0] * xq_[..., 0] + freqs[..., 1] * xq_[..., 1]
+    xk_ = freqs[..., 0] * xk_[..., 0] + freqs[..., 1] * xk_[..., 1]
+    return xq_.reshape(*xq.shape).to(xq.dtype), xk_.reshape(*xk.shape).to(xk.dtype)
+
+
+def attention(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    mask: Tensor | None = None,
+    scale: float | None = None,
+    gqa: bool = False,
+) -> Tensor:
+    with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
+        x = F.scaled_dot_product_attention(
+            q, k, v, attn_mask=mask, scale=scale, enable_gqa=gqa
+        )
+    return rearrange(x, "B H L D -> B L (H D)")
+
+
+def _mask(mask: Tensor) -> Tensor:
+    """Expand a (B, L) key-padding mask into a (B, 1, L, L) attention mask."""
+    return mask.unsqueeze(1).unsqueeze(2) * mask.unsqueeze(1).unsqueeze(3)
+
+
+def temb(
+    t: Tensor,
+    dim: int,
+    period: float = 1e4,
+    tfactor: float = 1e3,
+    device: torch.device = None,
+    dtype: torch.dtype = None,
+) -> Tensor:
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(period)
+        * torch.arange(half, dtype=torch.float32, device=device)
+        / half
+    )
+    # t: (B,) -> args: (B, 1, half), so the embedding broadcasts as a per-sample vec.
+    args = (t.float() * tfactor)[:, None, None] * freqs
+    sin, cos = torch.sin(args), torch.cos(args)
+    return torch.cat((cos, sin), dim=-1).to(dtype=dtype)
+
+
+@dataclass
+class SingleMMDiTConfig:
+    features: int
+    tdim: int
+    txtdim: int
+    heads: int
+    multiplier: int
+    layers: int
+    patch: int
+    channels: int
+    bias: bool = False
+    theta: float = 1e3
+    kvheads: int | None = None
+    txtlayers: int = 1
+    txtheads: int = 20
+    txtkvheads: int = 20
+
+
+class SimpleModulation(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.lin = torch.nn.Parameter(torch.zeros(2, dim))
+        self.multiplier = 2
+
+    # vec (b d)
+    def forward(self, vec: Tensor):
+        out = vec + rearrange(self.lin, "two d -> 1 two d")
+        scale, shift = out.chunk(self.multiplier, dim=1)
+        return scale, shift
+
+
+class DoubleSharedModulation(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.lin = torch.nn.Parameter(torch.zeros(6 * dim))
+
+    # vec (b (6 d))
+    def forward(self, vec: Tensor):
+        out = vec + self.lin
+        prescale, preshift, pregate, postscale, postshift, postgate = out.chunk(
+            6, dim=-1
+        )
+        return prescale, preshift, pregate, postscale, postshift, postgate
+
+
+class PositionalEncoding(torch.nn.Module):
+    def __init__(self, dim, axdims: list[int], theta: float = 1e2, ntk: float = 1.0):
+        super().__init__()
+        self.axdims = axdims  # how to split the head dimension across the position axes
+        self.theta = theta
+        self.ntk = ntk
+
+    @torch.compile(fullgraph=True)
+    def forward(self, pos: Tensor) -> Tensor:
+        return torch.cat(
+            [
+                rope(pos[..., i], d, self.theta, self.ntk)
+                for i, d in enumerate(self.axdims)
+            ],
+            dim=-3,
+        )
+
+
+class QKNorm(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.qnorm = RMSNorm(dim)
+        self.knorm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        return self.qnorm(q), self.knorm(k), v
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, features: int, eps: float = 1e-05, device: torch.device = None):
+        super().__init__()
+        self.features = features
+        self.eps = eps
+        self.scale = torch.nn.Parameter(
+            torch.zeros(features, device=device, dtype=torch.float32)
+        )
+
+    @torch.compile(fullgraph=True)
+    def forward(self, x: Tensor) -> Tensor:
+        t, dtype = x.float(), x.dtype
+        t = F.rms_norm(
+            t, (self.features,), eps=self.eps, weight=(self.scale.float() + 1.0)
+        )
+        return t.to(dtype)
+
+
+class SwiGLU(torch.nn.Module):
+    def __init__(
+        self, features: int, multiplier: int, bias: bool = False, multiple: int = 128
+    ):
+        super().__init__()
+
+        mlpdim = int(2 * features / 3) * multiplier
+        mlpdim = multiple * ((mlpdim + multiple - 1) // multiple)
+
+        self.gate = torch.nn.Linear(features, mlpdim, bias=bias)
+        self.up = torch.nn.Linear(features, mlpdim, bias=bias)
+        self.down = torch.nn.Linear(mlpdim, features, bias=bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class Attention(torch.nn.Module):
+    def __init__(self, dim: int, heads: int, kvheads: int = None, bias: bool = False):
+        super().__init__()
+        self.heads = heads
+        self.kvheads = kvheads if kvheads is not None else heads
+        self.headdim = dim // self.heads
+
+        self.wq = torch.nn.Linear(dim, self.headdim * self.heads, bias=bias)
+        self.wk = torch.nn.Linear(dim, self.headdim * self.kvheads, bias=bias)
+        self.wv = torch.nn.Linear(dim, self.headdim * self.kvheads, bias=bias)
+        self.gate = torch.nn.Linear(dim, dim, bias=bias)
+        self.qknorm = QKNorm(self.headdim)
+        self.gqa = self.heads != self.kvheads
+        self.wo = torch.nn.Linear(dim, dim, bias=bias)
+
+    def forward(
+        self, qkv: Tensor, freqs: Tensor | None = None, mask: Tensor | None = None
+    ) -> Tensor:
+        q, k, v, gate = self.wq(qkv), self.wk(qkv), self.wv(qkv), self.gate(qkv)
+
+        q, k, v = (
+            rearrange(q, "B L (H D) -> B H L D", H=self.heads),
+            rearrange(k, "B L (H D) -> B H L D", H=self.kvheads),
+            rearrange(v, "B L (H D) -> B H L D", H=self.kvheads),
+        )
+
+        q, k, v = self.qknorm(q, k, v)
+        if freqs is not None:
+            q, k = ropeapply(q, k, freqs)
+        out = self.wo(attention(q, k, v, mask=mask, gqa=self.gqa) * F.sigmoid(gate))
+
+        return out
+
+
+class LastLayer(torch.nn.Module):
+    def __init__(self, features: int, patch: int, channels: int):
+        super().__init__()
+        self.norm = RMSNorm(features)
+        self.linear = torch.nn.Linear(features, patch * patch * channels, bias=True)
+        self.modulation = SimpleModulation(features)
+
+    @torch.compile(fullgraph=True)
+    def forward(self, x: Tensor, tvec: Tensor) -> Tensor:
+        scale, shift = self.modulation(tvec)
+        x = (1 + scale) * self.norm(x) + shift
+        x = self.linear(x)
+        return x
+
+
+class TextFusionBlock(torch.nn.Module):
+    def __init__(
+        self,
+        features: int,
+        heads: int,
+        multiplier: int,
+        bias: bool = False,
+        kvheads: int = None,
+    ):
+        super().__init__()
+        self.prenorm = RMSNorm(features)
+        self.postnorm = RMSNorm(features)
+        self.attn = Attention(dim=features, heads=heads, bias=bias, kvheads=kvheads)
+        self.mlp = SwiGLU(features, multiplier, bias)
+
+    def forward(self, x: Tensor, mask: Tensor | None = None) -> Tensor:
+        x = x + self.attn(self.prenorm(x), mask=mask)
+        x = x + self.mlp(self.postnorm(x))
+
+        return x
+
+
+class TextFusionTransformer(torch.nn.Module):
+    # num_txt_layers is the number of selected encoder hidden-state layers fed in
+    # (projected down to 1), NOT the transformer depth — that's fixed at 2 + 2 blocks.
+    def __init__(
+        self,
+        num_txt_layers: int,
+        txt_dim: int,
+        heads: int,
+        multiplier: int,
+        bias: bool = False,
+        kvheads: int = None,
+    ):
+        super().__init__()
+        self.layerwise_blocks = torch.nn.ModuleList(
+            [
+                TextFusionBlock(txt_dim, heads, multiplier, bias, kvheads)
+                for _ in range(2)
+            ]
+        )
+        self.projector = torch.nn.Linear(num_txt_layers, 1, bias=False)
+        self.refiner_blocks = torch.nn.ModuleList(
+            [
+                TextFusionBlock(txt_dim, heads, multiplier, bias, kvheads)
+                for _ in range(2)
+            ]
+        )
+
+    def forward(self, x: Tensor, mask: Tensor | None = None) -> Tensor:
+        b, l, n, d = x.shape
+        x = x.reshape(b * l, n, d)
+        for block in self.layerwise_blocks:
+            x = block(x.contiguous(), mask=None)
+        x = rearrange(x, "(b l) n d -> b l d n", b=b, l=l)
+        x = self.projector(x)
+        x = x.squeeze(-1)
+
+        for block in self.refiner_blocks:
+            x = block(x, mask=mask)
+
+        return x
+
+
+class SingleStreamBlock(nn.Module):
+    def __init__(
+        self,
+        features: int,
+        heads: int,
+        multiplier: int,
+        bias: bool = False,
+        kvheads: int = None,
+    ):
+        super().__init__()
+        self.mod = DoubleSharedModulation(features)
+        self.prenorm = RMSNorm(features)
+        self.postnorm = RMSNorm(features)
+        self.attn = Attention(dim=features, heads=heads, bias=bias, kvheads=kvheads)
+        self.mlp = SwiGLU(features, multiplier, bias)
+
+    def forward(
+        self, x: Tensor, vec: Tensor, freqs: Tensor, mask: Tensor | None = None
+    ) -> Tensor:
+        prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec)
+        x = x + pregate * self.attn(
+            (1 + prescale) * self.prenorm(x) + preshift, freqs, mask
+        )
+        x = x + postgate * self.mlp((1 + postscale) * self.postnorm(x) + postshift)
+
+        return x
+
+
+class SingleStreamDiT(nn.Module):
+    def __init__(self, config: SingleMMDiTConfig):
+        super().__init__()
+        self.config = config
+
+        headdim = config.features // config.heads
+        axes = [
+            headdim - 12 * (headdim // 16),
+            6 * (headdim // 16),
+            6 * (headdim // 16),
+        ]
+        assert sum(axes) == headdim, f"sum(axes) = {sum(axes)}, headdim = {headdim}"
+        assert all(a % 2 == 0 for a in axes), f"axes = {axes}"
+
+        self.posemb = PositionalEncoding(
+            config.features, axes, theta=config.theta, ntk=1.0
+        )
+        self.first = nn.Linear(
+            config.channels * config.patch**2, config.features, bias=True
+        )
+
+        self.blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(
+                    config.features,
+                    config.heads,
+                    config.multiplier,
+                    config.bias,
+                    config.kvheads,
+                )
+                for _ in range(config.layers)
+            ]
+        )
+        self.tmlp = nn.Sequential(
+            nn.Linear(config.tdim, config.features),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(config.features, config.features),
+        )
+        self.txtfusion = TextFusionTransformer(
+            config.txtlayers,
+            config.txtdim,
+            config.txtheads,
+            config.multiplier,
+            config.bias,
+            config.txtkvheads,
+        )
+        self.txtmlp = nn.Sequential(
+            RMSNorm(config.txtdim),
+            nn.Linear(config.txtdim, config.features),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(config.features, config.features),
+        )
+        self.last = LastLayer(config.features, config.patch, config.channels)
+
+        self.tproj = nn.Sequential(
+            nn.GELU(approximate="tanh"), nn.Linear(config.features, config.features * 6)
+        )
+
+    def forward(
+        self,
+        img: Tensor,
+        context: Tensor,
+        t: Tensor,
+        pos: Tensor,
+        mask: Tensor | None = None,
+    ) -> Tensor:
+        img = self.first(img)
+        t = self.tmlp(temb(t, self.config.tdim, device=img.device, dtype=img.dtype))
+        tvec = self.tproj(t)
+
+        txtmask = _mask(mask[:, : context.shape[1]])
+
+        context = self.txtfusion(context, mask=txtmask)
+        context = self.txtmlp(context)
+
+        txtlen, imglen = context.shape[1], img.shape[1]
+        combined = torch.cat((context, img), dim=1)
+
+        # Pad combined sequence to a multiple of 256 to stabilize compiled kernel shapes.
+        fulllen = combined.shape[1]
+        _padlen = (-fulllen) % 256
+        if _padlen > 0:
+            combined = F.pad(combined, (0, 0, 0, _padlen))
+            mask = F.pad(mask, (0, _padlen), value=False)
+            pos = F.pad(pos, (0, 0, 0, _padlen))
+
+        mask = _mask(mask)
+
+        freqs = self.posemb(pos)
+
+        for block in self.blocks:
+            combined = block(combined, tvec, freqs, mask)
+
+        final = self.last(combined, t)
+        output = final[:, txtlen : txtlen + imglen, :]
+
+        return output

pyproject.toml

@@ -0,0 +1,23 @@
+[project]
+name = "krea-2-oss"
+version = "0.1.0"
+description = "Minimal inference code for the Krea 2 (K2) diffusion model"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "torch>=2.9.0",
+    "transformers==4.57.1",
+    "safetensors>=0.4.0",
+    "pillow>=10.0.0",
+    "einops>=0.7.0",
+    "diffusers>=0.32.0",
+    "click>=8.4.1",
+]
+
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true
+
+[tool.uv.sources]
+torch = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }]

sampling.py

@@ -0,0 +1,146 @@
+"""Functional flow-matching sampler for the K2 MMDiT (no Scheduler class)."""
+
+import math
+
+import torch
+from einops import rearrange, repeat
+from PIL import Image
+
+
+def roundup(value, multiple, name):
+    """Round `value` up to the nearest multiple, logging when padding is applied."""
+    aligned = ((value + multiple - 1) // multiple) * multiple
+    if aligned != value:
+        print(
+            f"[sample] {name}={value} is not a multiple of {multiple}; padding to {aligned}"
+        )
+    return aligned
+
+
+def prepare(img, txtlen, patch, txtmask):
+    """Patchify the latent and build the combined text+image position / mask tensors.
+
+    Returns (img_tokens, pos, mask).
+    """
+    b, _, h, w = img.shape
+    h_, w_ = h // patch, w // patch
+    imgids = torch.zeros((h_, w_, 3), device=img.device)
+    imgids[..., 1] = torch.arange(h_, device=img.device)[:, None]
+    imgids[..., 2] = torch.arange(w_, device=img.device)[None, :]
+    imgpos = repeat(imgids, "h w three -> b (h w) three", b=b, three=3)
+    imgmask = torch.ones(b, h_ * w_, device=img.device, dtype=torch.bool)
+    img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch)
+
+    txtpos = torch.zeros(b, txtlen, 3, device=img.device)
+    mask = torch.cat((txtmask, imgmask), dim=1)
+    pos = torch.cat((txtpos, imgpos), dim=1)
+    return img, pos, mask
+
+
+def timesteps(seq_len, steps, x1, x2, y1=0.5, y2=1.15, sigma=1.0, mu=None):
+    """Resolution-aware flow-matching timestep schedule (t: 1 -> 0).
+
+    `mu` is interpolated linearly in image-sequence length between (x1,y1) and
+    (x2,y2), then used to time-shift a uniform 1->0 grid. Pass an explicit `mu`
+    to pin a constant shift regardless of resolution (used by the distilled
+    checkpoint, which was trained at a fixed mu=1.15).
+    """
+    ts = torch.linspace(1, 0, steps + 1)
+    if mu is None:
+        slope = (y2 - y1) / (x2 - x1)
+        mu = slope * seq_len + (y1 - slope * x1)
+    ts = math.exp(mu) / (math.exp(mu) + (1.0 / ts - 1.0) ** sigma)
+    return ts.tolist()
+
+
+@torch.no_grad()
+def sample(
+    model,
+    ae,
+    encoder,
+    prompts,
+    *,
+    negative_prompts=None,
+    device="cuda",
+    dtype=torch.bfloat16,
+    width=1024,
+    height=1024,
+    steps=28,
+    guidance=4.5,
+    seed=0,
+    minres=256,
+    maxres=1280,
+    y1=0.5,
+    y2=1.15,
+    mu=None,
+):
+    """End-to-end text-to-image sampling: encode -> euler+CFG denoise -> decode."""
+    patch = model.config.patch
+
+    # The latent grid (dim // ae.compression) is patchified in `patch`-sized blocks,
+    # so width/height must be multiples of ae.compression * patch. Pad up otherwise.
+    align = ae.compression * patch
+    width, height = roundup(width, align, "width"), roundup(height, align, "height")
+
+    n = len(prompts)
+    cfg = guidance > 0
+    if negative_prompts is None:
+        negative_prompts = [""] * n
+
+    # Per-prompt seeded gaussian latent noise.
+    noise = torch.cat(
+        [
+            torch.randn(
+                1,
+                ae.channels,
+                height // ae.compression,
+                width // ae.compression,
+                device=device,
+                dtype=dtype,
+                generator=torch.Generator(device=device).manual_seed(seed + i),
+            )
+            for i in range(n)
+        ],
+        dim=0,
+    )
+
+    # Positive (conditional) text conditioning.
+    txt, txtmask = encoder(prompts)
+    x, pos, mask = prepare(noise, txt.shape[1], patch, txtmask)
+
+    # The unconditional branch is only used for CFG; skip encoding/prep entirely
+    # when guidance is disabled.
+    if cfg:
+        untxt, untxtmask = encoder(negative_prompts)
+        _, unpos, unmask = prepare(noise, untxt.shape[1], patch, untxtmask)
+
+    # min_res/max_res define the (x1,y1)-(x2,y2) interpolation endpoints for `mu`.
+    x1 = (minres // (ae.compression * patch)) ** 2
+    x2 = (maxres // (ae.compression * patch)) ** 2
+    ts = timesteps(x.shape[1], steps, x1, x2, y1=y1, y2=y2, mu=mu)
+
+    # Euler integration of the flow ODE with CFG.
+    img = x
+    for tcurr, tprev in zip(ts[:-1], ts[1:]):
+        t = torch.full((len(img),), tcurr, dtype=img.dtype, device=img.device)
+        cond = model(img=img, context=txt, t=t, pos=pos, mask=mask)
+        if cfg:
+            uncond = model(img=img, context=untxt, t=t, pos=unpos, mask=unmask)
+            v = cond + guidance * (cond - uncond)
+        else:
+            v = cond
+        img = img + (tprev - tcurr) * v
+
+    # Unpatchify back to a latent and decode to pixels.
+    img = rearrange(
+        img,
+        "b (h w) (c ph pw) -> b c (h ph) (w pw)",
+        ph=patch,
+        pw=patch,
+        h=height // (ae.compression * patch),
+        w=width // (ae.compression * patch),
+    )
+    img = ae.decode(img.to(torch.bfloat16))
+    img = img.clamp(-1, 1) * 0.5 + 0.5
+    img = rearrange(img * 255.0, "b c h w -> b h w c").cpu().byte().numpy()
+    return [Image.fromarray(img[i]) for i in range(len(img))]