inline-core 1.2.41__tar.gz → 1.2.52__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (184) hide show
  1. {inline_core-1.2.41 → inline_core-1.2.52}/PKG-INFO +40 -16
  2. {inline_core-1.2.41 → inline_core-1.2.52}/README.md +29 -13
  3. {inline_core-1.2.41 → inline_core-1.2.52}/pyproject.toml +19 -5
  4. inline_core-1.2.52/src/inline_core/models/checkpoint.py +87 -0
  5. inline_core-1.2.52/src/inline_core/models/krea2/__init__.py +9 -0
  6. inline_core-1.2.52/src/inline_core/models/krea2/convert.py +138 -0
  7. inline_core-1.2.52/src/inline_core/models/krea2/img2img.py +89 -0
  8. inline_core-1.2.52/src/inline_core/models/krea2/provider.py +57 -0
  9. inline_core-1.2.52/src/inline_core/models/krea2/requirements.py +180 -0
  10. inline_core-1.2.52/src/inline_core/models/krea2/runner.py +407 -0
  11. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/models/loaders.py +334 -22
  12. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/models/lora.py +78 -28
  13. inline_core-1.2.52/src/inline_core/models/pipeline_runtime.py +522 -0
  14. inline_core-1.2.52/src/inline_core/models/zimage/runner.py +465 -0
  15. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/server/app.py +24 -0
  16. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/server/bootstrap.py +10 -0
  17. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/handlers.py +6 -2
  18. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/training.py +30 -3
  19. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/training_store.py +6 -0
  20. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/training/__main__.py +8 -1
  21. inline_core-1.2.52/src/inline_core/training/arch.py +184 -0
  22. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/training/caption.py +39 -6
  23. inline_core-1.2.52/src/inline_core/training/dataset.py +158 -0
  24. inline_core-1.2.52/src/inline_core/training/models.py +319 -0
  25. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/training/protocol.py +8 -1
  26. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/training/trainer.py +81 -56
  27. inline_core-1.2.52/tests/test_checkpoint.py +79 -0
  28. inline_core-1.2.52/tests/test_krea2_convert.py +157 -0
  29. inline_core-1.2.52/tests/test_krea2_requirements.py +115 -0
  30. inline_core-1.2.52/tests/test_krea2_runner.py +163 -0
  31. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_loaders.py +15 -3
  32. inline_core-1.2.52/tests/test_lora_download.py +75 -0
  33. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_studio_rpc.py +2 -0
  34. inline_core-1.2.52/tests/test_training_arch.py +113 -0
  35. inline_core-1.2.52/tests/test_training_dataset.py +103 -0
  36. inline_core-1.2.52/tests/test_training_models.py +135 -0
  37. inline_core-1.2.52/tests/test_training_resolve.py +90 -0
  38. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_zimage_runner.py +17 -16
  39. {inline_core-1.2.41 → inline_core-1.2.52}/uv.lock +54 -54
  40. inline_core-1.2.52/webui.bat +225 -0
  41. inline_core-1.2.41/src/inline_core/models/zimage/runner.py +0 -985
  42. inline_core-1.2.41/src/inline_core/training/dataset.py +0 -100
  43. inline_core-1.2.41/src/inline_core/training/models.py +0 -125
  44. {inline_core-1.2.41 → inline_core-1.2.52}/.gitignore +0 -0
  45. {inline_core-1.2.41 → inline_core-1.2.52}/.python-version +0 -0
  46. {inline_core-1.2.41 → inline_core-1.2.52}/CLAUDE.md +0 -0
  47. {inline_core-1.2.41 → inline_core-1.2.52}/main.py +0 -0
  48. {inline_core-1.2.41 → inline_core-1.2.52}/scripts/reference.py +0 -0
  49. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/__init__.py +0 -0
  50. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/components/__init__.py +0 -0
  51. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/components/conditioning.py +0 -0
  52. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/components/interfaces.py +0 -0
  53. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/config.py +0 -0
  54. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/device/__init__.py +0 -0
  55. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/device/auto.py +0 -0
  56. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/device/detect.py +0 -0
  57. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/device/memory.py +0 -0
  58. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/device/policy.py +0 -0
  59. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/device/types.py +0 -0
  60. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/errors.py +0 -0
  61. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/__init__.py +0 -0
  62. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/api.py +0 -0
  63. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/constraints.py +0 -0
  64. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/fetch.py +0 -0
  65. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/handlers.py +0 -0
  66. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/importer.py +0 -0
  67. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/install.py +0 -0
  68. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/loader.py +0 -0
  69. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/manifest.py +0 -0
  70. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/models.py +0 -0
  71. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/paths.py +0 -0
  72. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/resolve.py +0 -0
  73. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/scanner.py +0 -0
  74. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/state.py +0 -0
  75. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/extensions/tools.py +0 -0
  76. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/ffmpeg.py +0 -0
  77. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/graph/__init__.py +0 -0
  78. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/graph/cache.py +0 -0
  79. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/graph/descriptor.py +0 -0
  80. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/graph/executor.py +0 -0
  81. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/graph/loader_runners.py +0 -0
  82. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/graph/primitives.py +0 -0
  83. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/graph/registry.py +0 -0
  84. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/graph/runners.py +0 -0
  85. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/graph/schema.py +0 -0
  86. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/graph/topo.py +0 -0
  87. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/graph/validate.py +0 -0
  88. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/media.py +0 -0
  89. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/models/__init__.py +0 -0
  90. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/models/catalog.py +0 -0
  91. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/models/requirements.py +0 -0
  92. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/models/sampling.py +0 -0
  93. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/models/zimage/__init__.py +0 -0
  94. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/models/zimage/primitives.py +0 -0
  95. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/models/zimage/provider.py +0 -0
  96. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/models/zimage/requirements.py +0 -0
  97. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/parallel/__init__.py +0 -0
  98. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/parallel/config.py +0 -0
  99. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/parallel/group.py +0 -0
  100. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/parallel/launch.py +0 -0
  101. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/parallel/protocol.py +0 -0
  102. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/parallel/registry.py +0 -0
  103. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/parallel/worker.py +0 -0
  104. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/runtime/__init__.py +0 -0
  105. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/runtime/context.py +0 -0
  106. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/runtime/file_store.py +0 -0
  107. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/runtime/progress.py +0 -0
  108. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/runtime/run.py +0 -0
  109. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/runtime/store.py +0 -0
  110. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/sampling/__init__.py +0 -0
  111. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/sampling/batch.py +0 -0
  112. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/server/__init__.py +0 -0
  113. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/server/__main__.py +0 -0
  114. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/server/assets.py +0 -0
  115. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/server/frontend.py +0 -0
  116. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/server/manager.py +0 -0
  117. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/server/rpc.py +0 -0
  118. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/server/run_store.py +0 -0
  119. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/server/serialize.py +0 -0
  120. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/__init__.py +0 -0
  121. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/assets.py +0 -0
  122. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/config.py +0 -0
  123. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/fal.py +0 -0
  124. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/frames.py +0 -0
  125. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/generation.py +0 -0
  126. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/graph_build.py +0 -0
  127. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/models.py +0 -0
  128. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/moodboard.py +0 -0
  129. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/peaks.py +0 -0
  130. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/schema.py +0 -0
  131. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/store.py +0 -0
  132. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/system_stats.py +0 -0
  133. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/timeline/__init__.py +0 -0
  134. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/timeline/compose.py +0 -0
  135. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/timeline/ffmpeg.py +0 -0
  136. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/timeline/render.py +0 -0
  137. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/studio/timeline/resolve.py +0 -0
  138. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/takes.py +0 -0
  139. {inline_core-1.2.41 → inline_core-1.2.52}/src/inline_core/training/__init__.py +0 -0
  140. {inline_core-1.2.41 → inline_core-1.2.52}/tests/helpers.py +0 -0
  141. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_cache.py +0 -0
  142. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_catalog.py +0 -0
  143. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_config.py +0 -0
  144. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_device_detect.py +0 -0
  145. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_executor.py +0 -0
  146. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_extension_api.py +0 -0
  147. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_extension_install.py +0 -0
  148. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_extension_manifest.py +0 -0
  149. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_extension_resolve.py +0 -0
  150. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_extension_scanner.py +0 -0
  151. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_extension_spine.py +0 -0
  152. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_extension_state.py +0 -0
  153. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_file_store.py +0 -0
  154. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_frontend_serving.py +0 -0
  155. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_hidden_nodes.py +0 -0
  156. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_loader_runners.py +0 -0
  157. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_lora.py +0 -0
  158. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_memory_policy.py +0 -0
  159. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_model_requirements.py +0 -0
  160. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_parallel_group.py +0 -0
  161. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_primitives.py +0 -0
  162. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_rpc_bridge.py +0 -0
  163. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_run_store.py +0 -0
  164. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_sampling.py +0 -0
  165. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_schema.py +0 -0
  166. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_server.py +0 -0
  167. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_studio_assets.py +0 -0
  168. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_studio_fal.py +0 -0
  169. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_studio_frames.py +0 -0
  170. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_studio_generation.py +0 -0
  171. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_studio_models.py +0 -0
  172. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_studio_moodboard.py +0 -0
  173. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_studio_peaks.py +0 -0
  174. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_studio_schema.py +0 -0
  175. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_studio_store.py +0 -0
  176. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_studio_timeline.py +0 -0
  177. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_studio_training.py +0 -0
  178. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_take_bytes.py +0 -0
  179. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_topo.py +0 -0
  180. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_validate.py +0 -0
  181. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_xfuser_sampler.py +0 -0
  182. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_zimage_primitives.py +0 -0
  183. {inline_core-1.2.41 → inline_core-1.2.52}/tests/test_zimage_resolve.py +0 -0
  184. {inline_core-1.2.41 → inline_core-1.2.52}/webui.sh +0 -0
@@ -1,18 +1,26 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: inline-core
3
- Version: 1.2.41
3
+ Version: 1.2.52
4
4
  Summary: The generation engine behind Inline Studio.
5
+ License-Expression: GPL-3.0-or-later
5
6
  Requires-Python: >=3.11
6
7
  Requires-Dist: numpy>=1.26
7
8
  Requires-Dist: psutil>=5.9
8
9
  Provides-Extra: all
9
10
  Requires-Dist: accelerate>=0.30; extra == 'all'
10
- Requires-Dist: diffusers>=0.36; extra == 'all'
11
+ Requires-Dist: bitsandbytes>=0.43; (platform_system != 'Darwin') and extra == 'all'
12
+ Requires-Dist: diffusers>=0.39; extra == 'all'
13
+ Requires-Dist: einops>=0.7; extra == 'all'
11
14
  Requires-Dist: fastapi>=0.110; extra == 'all'
12
15
  Requires-Dist: huggingface-hub>=0.23; extra == 'all'
13
16
  Requires-Dist: imageio-ffmpeg>=0.4; extra == 'all'
17
+ Requires-Dist: nvidia-ml-py>=12; extra == 'all'
18
+ Requires-Dist: peft>=0.11; extra == 'all'
19
+ Requires-Dist: pillow>=10; extra == 'all'
20
+ Requires-Dist: psutil>=5.9; extra == 'all'
14
21
  Requires-Dist: safetensors>=0.4; extra == 'all'
15
22
  Requires-Dist: scipy>=1.11; extra == 'all'
23
+ Requires-Dist: timm>=1.0; extra == 'all'
16
24
  Requires-Dist: torch>=2.2; extra == 'all'
17
25
  Requires-Dist: torchao>=0.14; extra == 'all'
18
26
  Requires-Dist: transformers>=4.44; extra == 'all'
@@ -27,7 +35,7 @@ Requires-Dist: nvidia-ml-py>=12; extra == 'parallel'
27
35
  Requires-Dist: xfuser>=0.4; extra == 'parallel'
28
36
  Provides-Extra: runtime
29
37
  Requires-Dist: accelerate>=0.30; extra == 'runtime'
30
- Requires-Dist: diffusers>=0.36; extra == 'runtime'
38
+ Requires-Dist: diffusers>=0.39; extra == 'runtime'
31
39
  Requires-Dist: huggingface-hub>=0.23; extra == 'runtime'
32
40
  Requires-Dist: safetensors>=0.4; extra == 'runtime'
33
41
  Requires-Dist: scipy>=1.11; extra == 'runtime'
@@ -55,8 +63,9 @@ The generation engine behind Inline. It takes a typed node graph (JSON) and retu
55
63
  low-VRAM laptops up to multi-GPU machines that split a single image's sampling across GPUs (via
56
64
  xDiT). It is Inline Studio's built-in render backend.
57
65
 
58
- First model: Z-Image (Alibaba Tongyi), a 6B rectified-flow diffusion transformer (and a model xDiT
59
- already supports, so the multi-GPU split works on it from the start).
66
+ Models today: **Z-Image** (Alibaba Tongyi), a 6B rectified-flow diffusion transformer and the model
67
+ xDiT already supports, so the multi-GPU split works on it from the start; and **Krea 2** (Krea AI),
68
+ a 12.9B single-stream MMDiT released as an undistilled RAW base plus an 8-step distilled Turbo.
60
69
 
61
70
  > Status: early, and running end to end against a stub engine. In place and tested: the graph engine,
62
71
  > the typed `/v1` HTTP + websocket API (durable runs, streamed progress, coalescing), the model-dir
@@ -102,9 +111,13 @@ by category:
102
111
 
103
112
  ```
104
113
  models/
105
- diffusion_models/ z_image_turbo_bf16.safetensors <- the one file you need for Z-Image
106
- vae/ ae.safetensors <- optional (see below)
107
- text_encoders/ qwen/ (an HF-format folder: config + tokenizer + weights) <- optional
114
+ diffusion_models/ z_image_turbo_bf16.safetensors <- Z-Image
115
+ krea2_turbo_bf16.safetensors <- Krea 2 Turbo (generate)
116
+ krea2_raw_bf16.safetensors <- Krea 2 RAW (fine-tune)
117
+ vae/ ae.safetensors <- Z-Image
118
+ qwen_image_vae_diffusers.safetensors <- Krea 2
119
+ text_encoders/ qwen_3_4b.safetensors <- Z-Image
120
+ qwen3vl_4b_bf16.safetensors <- Krea 2
108
121
  loras/ controlnet/ checkpoints/ ...
109
122
  ```
110
123
 
@@ -123,8 +136,17 @@ paths:
123
136
  `models/`** (never the hidden HF cache) with visible progress.
124
137
 
125
138
  Override the diffusion source with `INLINE_ZIMAGE_MODEL` (a file or a diffusers dir), and the supporting
126
- components with `INLINE_ZIMAGE_VAE` / `INLINE_ZIMAGE_TEXT_ENCODER`. The engine scans the models dir on
127
- start; a node's model pickers list what is present.
139
+ components with `INLINE_ZIMAGE_VAE` / `INLINE_ZIMAGE_TEXT_ENCODER`; Krea 2 has the same three under
140
+ `INLINE_KREA2_*`. The engine scans the models dir on start; a node's model pickers list what is present.
141
+
142
+ **Krea 2 loads from the same ComfyUI single files**, with two constraints worth knowing. Only the
143
+ **bf16** builds in [`Comfy-Org/Krea-2`](https://huggingface.co/Comfy-Org/Krea-2) are loadable: the
144
+ fp8/int8/nvfp4 variants carry ComfyUI-specific scale tensors, and the loader refuses them by name
145
+ rather than failing mid-load. And the **VAE must be the diffusers-format** Qwen-Image file (the
146
+ node's popup fetches it from `Qwen/Qwen-Image`), because ComfyUI's copy of the same weights uses a
147
+ module layout diffusers has no converter for. `Krea2Transformer2DModel` has no `from_single_file`, so
148
+ the transformer checkpoint is renamed into diffusers naming on load (`models/krea2/convert.py`) and
149
+ streamed tensor by tensor onto the GPU.
128
150
 
129
151
  ## Nodes
130
152
 
@@ -132,8 +154,9 @@ start; a node's model pickers list what is present.
132
154
  renders any node generically - adding a node type needs no UI release.
133
155
 
134
156
  **High-level model nodes are what the user sees.** Generation is one-click: you drop a single
135
- **Z-Image Turbo** node, wire a Prompt into it, and hit Run. The node hooks up the diffusion model, VAE,
136
- and text-encoder behind the scenes - no loader/sampler wiring.
157
+ **Z-Image Turbo**, **Krea 2 Turbo** or **Krea 2 RAW** node, wire a Prompt into it, and hit Run. The
158
+ node hooks up the diffusion model, VAE, and text-encoder behind the scenes - no loader/sampler
159
+ wiring.
137
160
 
138
161
  Underneath, a **low-level primitive vocabulary** exists (`load/diffusion-model`, `load/vae`,
139
162
  `load/text-encoder`, `encode/text`, `latent/empty`, `sample`, `vae/decode`, `vae/encode`) - the
@@ -215,17 +238,18 @@ offload machinery, `--vram-budget GB` / `INLINE_VRAM_BUDGET_GB` to cap the assum
215
238
 
216
239
  ## Run
217
240
 
218
- The easy path is `webui.sh`, which maps friendly flags onto the engine's `INLINE_*` env knobs:
241
+ The easy path is `webui.sh` (macOS/Linux) or `webui.bat` (Windows), which maps friendly flags onto the
242
+ engine's `INLINE_*` env knobs. Windows must use `webui.bat` - `webui.sh` is a bash script:
219
243
 
220
244
  ```
221
- ./webui.sh # loopback, port 8848
245
+ ./webui.sh # loopback, port 8848 (Windows: .\webui.bat)
222
246
  ./webui.sh --listen --port 9000 # bind all interfaces on 9000
223
247
  ./webui.sh --multi-gpu # split one image across GPUs (auto with 2+ GPUs)
224
248
  ./webui.sh --lowvram # tight-VRAM profile
225
- ./webui.sh --install --extra zimage # set up ./.venv with the Z-Image runtime, then exit
249
+ ./webui.sh --install --extra runtime # set up ./.venv with the model runtime, then exit
226
250
  ```
227
251
 
228
- `./webui.sh --help` lists every flag (networking, multi-GPU, device/memory profile, paths). The same
252
+ `./webui.sh --help` (or `.\webui.bat --help`) lists every flag (networking, multi-GPU, device/memory profile, paths). The same
229
253
  flags are available on the Python entrypoint - `python main.py --help` - which is the dev path (it also
230
254
  takes `--front-end-root DIR` to serve a local SPA build). Or run the server module directly:
231
255
 
@@ -5,8 +5,9 @@ The generation engine behind Inline. It takes a typed node graph (JSON) and retu
5
5
  low-VRAM laptops up to multi-GPU machines that split a single image's sampling across GPUs (via
6
6
  xDiT). It is Inline Studio's built-in render backend.
7
7
 
8
- First model: Z-Image (Alibaba Tongyi), a 6B rectified-flow diffusion transformer (and a model xDiT
9
- already supports, so the multi-GPU split works on it from the start).
8
+ Models today: **Z-Image** (Alibaba Tongyi), a 6B rectified-flow diffusion transformer and the model
9
+ xDiT already supports, so the multi-GPU split works on it from the start; and **Krea 2** (Krea AI),
10
+ a 12.9B single-stream MMDiT released as an undistilled RAW base plus an 8-step distilled Turbo.
10
11
 
11
12
  > Status: early, and running end to end against a stub engine. In place and tested: the graph engine,
12
13
  > the typed `/v1` HTTP + websocket API (durable runs, streamed progress, coalescing), the model-dir
@@ -52,9 +53,13 @@ by category:
52
53
 
53
54
  ```
54
55
  models/
55
- diffusion_models/ z_image_turbo_bf16.safetensors <- the one file you need for Z-Image
56
- vae/ ae.safetensors <- optional (see below)
57
- text_encoders/ qwen/ (an HF-format folder: config + tokenizer + weights) <- optional
56
+ diffusion_models/ z_image_turbo_bf16.safetensors <- Z-Image
57
+ krea2_turbo_bf16.safetensors <- Krea 2 Turbo (generate)
58
+ krea2_raw_bf16.safetensors <- Krea 2 RAW (fine-tune)
59
+ vae/ ae.safetensors <- Z-Image
60
+ qwen_image_vae_diffusers.safetensors <- Krea 2
61
+ text_encoders/ qwen_3_4b.safetensors <- Z-Image
62
+ qwen3vl_4b_bf16.safetensors <- Krea 2
58
63
  loras/ controlnet/ checkpoints/ ...
59
64
  ```
60
65
 
@@ -73,8 +78,17 @@ paths:
73
78
  `models/`** (never the hidden HF cache) with visible progress.
74
79
 
75
80
  Override the diffusion source with `INLINE_ZIMAGE_MODEL` (a file or a diffusers dir), and the supporting
76
- components with `INLINE_ZIMAGE_VAE` / `INLINE_ZIMAGE_TEXT_ENCODER`. The engine scans the models dir on
77
- start; a node's model pickers list what is present.
81
+ components with `INLINE_ZIMAGE_VAE` / `INLINE_ZIMAGE_TEXT_ENCODER`; Krea 2 has the same three under
82
+ `INLINE_KREA2_*`. The engine scans the models dir on start; a node's model pickers list what is present.
83
+
84
+ **Krea 2 loads from the same ComfyUI single files**, with two constraints worth knowing. Only the
85
+ **bf16** builds in [`Comfy-Org/Krea-2`](https://huggingface.co/Comfy-Org/Krea-2) are loadable: the
86
+ fp8/int8/nvfp4 variants carry ComfyUI-specific scale tensors, and the loader refuses them by name
87
+ rather than failing mid-load. And the **VAE must be the diffusers-format** Qwen-Image file (the
88
+ node's popup fetches it from `Qwen/Qwen-Image`), because ComfyUI's copy of the same weights uses a
89
+ module layout diffusers has no converter for. `Krea2Transformer2DModel` has no `from_single_file`, so
90
+ the transformer checkpoint is renamed into diffusers naming on load (`models/krea2/convert.py`) and
91
+ streamed tensor by tensor onto the GPU.
78
92
 
79
93
  ## Nodes
80
94
 
@@ -82,8 +96,9 @@ start; a node's model pickers list what is present.
82
96
  renders any node generically - adding a node type needs no UI release.
83
97
 
84
98
  **High-level model nodes are what the user sees.** Generation is one-click: you drop a single
85
- **Z-Image Turbo** node, wire a Prompt into it, and hit Run. The node hooks up the diffusion model, VAE,
86
- and text-encoder behind the scenes - no loader/sampler wiring.
99
+ **Z-Image Turbo**, **Krea 2 Turbo** or **Krea 2 RAW** node, wire a Prompt into it, and hit Run. The
100
+ node hooks up the diffusion model, VAE, and text-encoder behind the scenes - no loader/sampler
101
+ wiring.
87
102
 
88
103
  Underneath, a **low-level primitive vocabulary** exists (`load/diffusion-model`, `load/vae`,
89
104
  `load/text-encoder`, `encode/text`, `latent/empty`, `sample`, `vae/decode`, `vae/encode`) - the
@@ -165,17 +180,18 @@ offload machinery, `--vram-budget GB` / `INLINE_VRAM_BUDGET_GB` to cap the assum
165
180
 
166
181
  ## Run
167
182
 
168
- The easy path is `webui.sh`, which maps friendly flags onto the engine's `INLINE_*` env knobs:
183
+ The easy path is `webui.sh` (macOS/Linux) or `webui.bat` (Windows), which maps friendly flags onto the
184
+ engine's `INLINE_*` env knobs. Windows must use `webui.bat` - `webui.sh` is a bash script:
169
185
 
170
186
  ```
171
- ./webui.sh # loopback, port 8848
187
+ ./webui.sh # loopback, port 8848 (Windows: .\webui.bat)
172
188
  ./webui.sh --listen --port 9000 # bind all interfaces on 9000
173
189
  ./webui.sh --multi-gpu # split one image across GPUs (auto with 2+ GPUs)
174
190
  ./webui.sh --lowvram # tight-VRAM profile
175
- ./webui.sh --install --extra zimage # set up ./.venv with the Z-Image runtime, then exit
191
+ ./webui.sh --install --extra runtime # set up ./.venv with the model runtime, then exit
176
192
  ```
177
193
 
178
- `./webui.sh --help` lists every flag (networking, multi-GPU, device/memory profile, paths). The same
194
+ `./webui.sh --help` (or `.\webui.bat --help`) lists every flag (networking, multi-GPU, device/memory profile, paths). The same
179
195
  flags are available on the Python entrypoint - `python main.py --help` - which is the dev path (it also
180
196
  takes `--front-end-root DIR` to serve a local SPA build). Or run the server module directly:
181
197
 
@@ -1,9 +1,10 @@
1
1
  [project]
2
2
  # PyPI name; the import package is `inline_core` (src/inline_core).
3
3
  name = "inline-core"
4
- version = "1.2.41"
4
+ version = "1.2.52"
5
5
  description = "The generation engine behind Inline Studio."
6
6
  readme = "README.md"
7
+ license = "GPL-3.0-or-later"
7
8
  requires-python = ">=3.11"
8
9
  # Deliberately torch-free so the engine boots and serves source nodes with no ML stack installed.
9
10
  dependencies = [
@@ -16,8 +17,8 @@ dependencies = [
16
17
  runtime = [
17
18
  # A bare `torch` is CPU-only on Windows; see [tool.uv] below and requirements.txt.
18
19
  "torch>=2.2",
19
- # Z-Image needs diffusers >= 0.36 (ZImagePipeline).
20
- "diffusers>=0.36",
20
+ # Krea 2 needs diffusers >= 0.39 (Krea2Pipeline); Z-Image needs >= 0.36 (ZImagePipeline).
21
+ "diffusers>=0.39",
21
22
  "transformers>=4.44",
22
23
  "accelerate>=0.30",
23
24
  "safetensors>=0.4",
@@ -59,19 +60,32 @@ dev = [
59
60
  "pyright>=1.1",
60
61
  "httpx>=0.27",
61
62
  ]
62
- # Everything needed to run locally. Lists concrete packages, never `inline-core[...]` self-refs.
63
+ # Everything a local install needs in one extra: runtime + server + training (the LoRA trainer),
64
+ # so `--install --extra all` is a single "install everything" command. Multi-GPU (xfuser, the
65
+ # `parallel` extra) is left out on purpose - it is niche and can fail to build on many systems.
66
+ # Lists concrete packages, never `inline-core[...]` self-refs.
63
67
  all = [
68
+ # runtime
64
69
  "torch>=2.2",
65
- "diffusers>=0.36",
70
+ "diffusers>=0.39",
66
71
  "transformers>=4.44",
67
72
  "accelerate>=0.30",
68
73
  "safetensors>=0.4",
69
74
  "torchao>=0.14",
70
75
  "scipy>=1.11",
71
76
  "huggingface_hub>=0.23",
77
+ # server
72
78
  "fastapi>=0.110",
73
79
  "uvicorn[standard]>=0.29",
74
80
  "imageio-ffmpeg>=0.4",
81
+ # training (the Trainer tab)
82
+ "peft>=0.11",
83
+ "bitsandbytes>=0.43; platform_system != 'Darwin'",
84
+ "Pillow>=10",
85
+ "psutil>=5.9",
86
+ "nvidia-ml-py>=12",
87
+ "timm>=1.0",
88
+ "einops>=0.7",
75
89
  ]
76
90
 
77
91
  # PyPI ships CPU-only torch on Windows; pin the CUDA index there (pip: see requirements.txt).
@@ -0,0 +1,87 @@
1
+ """Read a safetensors checkpoint tensor by tensor, without mapping the whole file.
2
+
3
+ ``safetensors.safe_open`` maps the entire file at once, and Linux refuses a mapping larger than
4
+ physical RAM when there is no swap (the default heuristic overcommit mode). A 26GB Krea 2 checkpoint
5
+ is therefore unreadable on a 16GB machine - it fails with ``Cannot allocate memory`` before any GPU
6
+ work, no matter how small the model would be once quantized.
7
+
8
+ Reading each tensor's byte range instead keeps peak host RAM at one tensor and works for any file
9
+ size. Only used for the big single-file checkpoints; small files (LoRAs, VAEs) still go through
10
+ safetensors directly.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ import json
16
+ import struct
17
+ from pathlib import Path
18
+ from typing import Any
19
+
20
+ from ..errors import ComponentError
21
+
22
+ #: safetensors dtype names -> torch dtypes, resolved lazily so importing this module stays cheap.
23
+ _DTYPE_NAMES = {
24
+ "F64": "float64",
25
+ "F32": "float32",
26
+ "F16": "float16",
27
+ "BF16": "bfloat16",
28
+ "I64": "int64",
29
+ "I32": "int32",
30
+ "I16": "int16",
31
+ "I8": "int8",
32
+ "U8": "uint8",
33
+ "BOOL": "bool",
34
+ }
35
+
36
+
37
+ class CheckpointReader:
38
+ """Random access to one safetensors file, one tensor at a time."""
39
+
40
+ def __init__(self, path: str | Path) -> None:
41
+ self._path = Path(path)
42
+ try:
43
+ size = self._path.stat().st_size
44
+ with self._path.open("rb") as handle:
45
+ (header_len,) = struct.unpack("<Q", handle.read(8))
46
+ # Garbage bytes decode to an enormous length; reject it rather than trying to
47
+ # allocate it, which would raise MemoryError instead of a usable message.
48
+ if not 0 < header_len <= size - 8:
49
+ raise ValueError("header length is not plausible for this file")
50
+ header = json.loads(handle.read(header_len))
51
+ if not isinstance(header, dict):
52
+ raise ValueError("header is not a JSON object")
53
+ except (OSError, ValueError, struct.error) as error:
54
+ raise ComponentError(f"Could not read checkpoint {self._path.name}: {error}") from error
55
+ self._start = 8 + header_len
56
+ self._index: dict[str, Any] = {k: v for k, v in header.items() if k != "__metadata__"}
57
+ self.metadata: dict[str, Any] = header.get("__metadata__") or {}
58
+
59
+ def keys(self) -> list[str]:
60
+ return list(self._index)
61
+
62
+ def get_tensor(self, key: str, device: str | None = None) -> Any:
63
+ """One tensor, read straight from its byte range into a fresh buffer."""
64
+ import torch
65
+
66
+ entry = self._index[key]
67
+ dtype_name = _DTYPE_NAMES.get(entry["dtype"])
68
+ if dtype_name is None:
69
+ raise ComponentError(
70
+ f"Checkpoint {self._path.name} uses the unsupported dtype {entry['dtype']!r} "
71
+ f"for {key!r}."
72
+ )
73
+ start, end = entry["data_offsets"]
74
+ buffer = bytearray(end - start)
75
+ with self._path.open("rb") as handle:
76
+ handle.seek(self._start + start)
77
+ if handle.readinto(buffer) != len(buffer):
78
+ raise ComponentError(f"Checkpoint {self._path.name} is truncated at {key!r}.")
79
+ tensor = torch.frombuffer(buffer, dtype=getattr(torch, dtype_name))
80
+ tensor = tensor.reshape(entry["shape"]) if entry["shape"] else tensor.reshape(())
81
+ return tensor.to(device) if device else tensor
82
+
83
+ def __enter__(self) -> CheckpointReader:
84
+ return self
85
+
86
+ def __exit__(self, *_exc: object) -> None:
87
+ return None
@@ -0,0 +1,9 @@
1
+ """The Krea 2 (Krea AI) runtime: a diffusers-backed text-to-image / img2img runner.
2
+
3
+ Optional subpackage. `server.bootstrap` imports `register_krea2` best-effort, so a core install
4
+ without the ``runtime`` extra (torch + diffusers) still boots and serves the source nodes.
5
+ """
6
+
7
+ from .runner import register_krea2
8
+
9
+ __all__ = ["register_krea2"]
@@ -0,0 +1,138 @@
1
+ """Rename a Krea 2 checkpoint from the reference (ComfyUI) layout to the diffusers one.
2
+
3
+ ``Krea2Transformer2DModel`` has no ``from_single_file``, so the single ``.safetensors`` under
4
+ ``diffusion_models/`` cannot be loaded directly. The two layouts describe the same 430 tensors under
5
+ different names, so a pure rename plus one reshape is the whole conversion.
6
+
7
+ The same rename maps a LoRA's module path (``module_alias``), which matters because the two Krea 2
8
+ LoRA conventions disagree: the official Comfy-Org style LoRAs already use diffusers names, while
9
+ ostris' training adapter uses reference ones.
10
+
11
+ Torch-free apart from the tensors handed in, so it is cheap to import and easy to test.
12
+ """
13
+
14
+ from __future__ import annotations
15
+
16
+ import re
17
+ from collections.abc import Mapping
18
+ from typing import Any
19
+
20
+ from ...errors import ComponentError
21
+
22
+ #: Root-anchored renames, longest-first so ``txtmlp.0.scale`` wins over the ``txtmlp.1.`` prefix.
23
+ _ROOT: tuple[tuple[str, str], ...] = (
24
+ ("txtmlp.0.scale", "txt_in.norm.weight"),
25
+ ("txtmlp.1.", "txt_in.linear_1."),
26
+ ("txtmlp.3.", "txt_in.linear_2."),
27
+ ("tmlp.0.", "time_embed.linear_1."),
28
+ ("tmlp.2.", "time_embed.linear_2."),
29
+ ("tproj.1.", "time_mod_proj."),
30
+ ("first.", "img_in."),
31
+ ("last.linear.", "final_layer.linear."),
32
+ ("last.norm.scale", "final_layer.norm.weight"),
33
+ ("last.modulation.lin", "final_layer.scale_shift_table"),
34
+ ("txtfusion.", "text_fusion."),
35
+ ("blocks.", "transformer_blocks."),
36
+ )
37
+
38
+ #: Renames applied anywhere in the path - they name parts of a block, which appears at three depths
39
+ #: (the main stack and the two text-fusion stacks).
40
+ _INNER: tuple[tuple[str, str], ...] = (
41
+ (".attn.qknorm.qnorm.scale", ".attn.norm_q.weight"),
42
+ (".attn.qknorm.knorm.scale", ".attn.norm_k.weight"),
43
+ (".attn.wq.", ".attn.to_q."),
44
+ (".attn.wk.", ".attn.to_k."),
45
+ (".attn.wv.", ".attn.to_v."),
46
+ (".attn.wo.", ".attn.to_out.0."),
47
+ (".attn.gate.", ".attn.to_gate."),
48
+ (".mlp.up.", ".ff.up."),
49
+ (".mlp.gate.", ".ff.gate."),
50
+ (".mlp.down.", ".ff.down."),
51
+ (".prenorm.scale", ".norm1.weight"),
52
+ (".postnorm.scale", ".norm2.weight"),
53
+ )
54
+
55
+ #: Suffixes ComfyUI's quantized builds add on top of the 430 reference tensors.
56
+ _QUANT_MARKERS = (".weight_scale", ".weight_scale_2", ".comfy_quant")
57
+
58
+ _MOD_LIN = re.compile(r"\.mod\.lin$")
59
+
60
+
61
+ def convert_key(key: str) -> str:
62
+ """One reference key in the diffusers naming. An already-diffusers key passes through."""
63
+ for old, new in _ROOT:
64
+ if key.startswith(old):
65
+ key = new + key[len(old) :]
66
+ break
67
+ for old, new in _INNER:
68
+ key = key.replace(old, new)
69
+ return _MOD_LIN.sub(".scale_shift_table", key)
70
+
71
+
72
+ def module_alias(stem: str) -> str | None:
73
+ """A LoRA module path in diffusers naming, or None when the rename leaves it unchanged.
74
+
75
+ ``lora.py`` has already stripped the checkpoint prefix, so ``stem`` looks like
76
+ ``blocks.0.attn.wq``. The trailing ``.weight`` makes the suffix-anchored rules above fire."""
77
+ converted = convert_key(f"{stem}.weight").removesuffix(".weight")
78
+ return converted if converted != stem else None
79
+
80
+
81
+ def is_quantized_checkpoint(keys: Any) -> bool:
82
+ """Whether this is a ComfyUI fp8/int8/nvfp4 build rather than the bf16 one."""
83
+ return any(str(k).endswith(_QUANT_MARKERS) for k in keys)
84
+
85
+
86
+ def check_loadable(keys: Any, expected: set[str]) -> None:
87
+ """Refuse a checkpoint that would only partially load, before any tensor is read.
88
+
89
+ Key-only so the streaming loader can validate up front; a partial load would otherwise leave
90
+ random-initialised layers and produce quietly wrong images instead of an error."""
91
+ keys = list(keys)
92
+ if is_quantized_checkpoint(keys):
93
+ raise ComponentError(
94
+ "This is a ComfyUI quantized Krea 2 build (fp8 / int8 / nvfp4), which only ComfyUI can "
95
+ "read. Use krea2_raw_bf16.safetensors or krea2_turbo_bf16.safetensors - smart memory "
96
+ "quantizes it for your GPU on load."
97
+ )
98
+ converted = {convert_key(k): k for k in keys}
99
+ unknown = sorted(original for name, original in converted.items() if name not in expected)
100
+ if unknown:
101
+ raise ComponentError(
102
+ f"Krea 2 checkpoint has {len(unknown)} unrecognised tensors (e.g. "
103
+ f"{', '.join(unknown[:3])}). It is probably a different model or a quantized build."
104
+ )
105
+ missing = sorted(expected - set(converted))
106
+ if missing:
107
+ raise ComponentError(
108
+ f"Krea 2 checkpoint is missing {len(missing)} tensors (e.g. "
109
+ f"{', '.join(missing[:3])}). The file looks truncated or is not a full checkpoint."
110
+ )
111
+
112
+
113
+ def convert_state_dict(
114
+ state: Mapping[str, Any], shapes: Mapping[str, tuple[int, ...]]
115
+ ) -> dict[str, Any]:
116
+ """Rename ``state`` into ``shapes``' naming, reshaping the one entry whose layout differs
117
+ (``blocks.N.mod.lin`` is flat where ``scale_shift_table`` is 2-D)."""
118
+ check_loadable(state.keys(), set(shapes))
119
+ converted: dict[str, Any] = {}
120
+ for key, tensor in state.items():
121
+ name = convert_key(key)
122
+ want = tuple(shapes[name])
123
+ if tuple(tensor.shape) != want:
124
+ if tensor.numel() != _numel(want):
125
+ raise ComponentError(
126
+ f"Krea 2 checkpoint tensor {key!r} is {tuple(tensor.shape)}, but the model "
127
+ f"expects {want}. This file is not a Krea 2 transformer."
128
+ )
129
+ tensor = tensor.reshape(want)
130
+ converted[name] = tensor
131
+ return converted
132
+
133
+
134
+ def _numel(shape: tuple[int, ...]) -> int:
135
+ total = 1
136
+ for dim in shape:
137
+ total *= dim
138
+ return total
@@ -0,0 +1,89 @@
1
+ """Krea 2 img2img: start the schedule partway, from an encoded input image.
2
+
3
+ diffusers ships no ``Krea2Img2ImgPipeline``, but ``Krea2Pipeline.__call__`` already accepts
4
+ pre-noised ``latents`` and an explicit ``sigmas`` schedule, so img2img needs no fork of the denoise
5
+ loop - only the right latents and the right tail of the schedule.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ from typing import Any
11
+
12
+ import torch
13
+
14
+
15
+ def img2img_kwargs(
16
+ pipe: Any,
17
+ *,
18
+ image: Any,
19
+ strength: float,
20
+ steps: int,
21
+ width: int,
22
+ height: int,
23
+ generator: Any,
24
+ device: str,
25
+ ) -> dict[str, Any]:
26
+ """Start the schedule partway, from the input image, using the pipeline's own ``latents`` and
27
+ ``sigmas`` inputs - diffusers ships no ``Krea2Img2ImgPipeline`` and forking its denoise loop
28
+ would be a maintenance liability.
29
+
30
+ The sigma tail is handed over **unshifted**: ``set_timesteps`` applies the same monotone
31
+ exponential shift the pipeline would, so a shifted subset and a subset of the shifted schedule
32
+ are the same thing. The latents are noised at the *shifted* start sigma so the noise level the
33
+ model sees matches the timestep it is given."""
34
+ import numpy as np
35
+ from diffusers.pipelines.krea2.pipeline_krea2 import calculate_shift
36
+
37
+ raw = np.linspace(1.0, 1 / steps, steps)
38
+ start = min(steps - 1, max(0, int(round((1.0 - strength) * steps))))
39
+
40
+ vae_scale = pipe.vae_scale_factor * pipe.patch_size
41
+ grid_h, grid_w = height // vae_scale, width // vae_scale
42
+ if pipe.config.is_distilled:
43
+ mu = 1.15
44
+ else:
45
+ config = pipe.scheduler.config
46
+ mu = calculate_shift(
47
+ grid_h * grid_w,
48
+ config.get("base_image_seq_len", 256),
49
+ config.get("max_image_seq_len", 6400),
50
+ config.get("base_shift", 0.5),
51
+ config.get("max_shift", 1.15),
52
+ )
53
+ schedule = pipe.scheduler.__class__.from_config(pipe.scheduler.config)
54
+ schedule.set_timesteps(sigmas=raw, mu=mu, device=device)
55
+ sigma = schedule.sigmas[start].to(device)
56
+
57
+ latents = _encode_image(pipe, image, width, height, device, generator)
58
+ noise = torch.randn(
59
+ latents.shape, generator=generator, device=latents.device, dtype=latents.dtype
60
+ )
61
+ return {"latents": (1.0 - sigma) * latents + sigma * noise, "sigmas": raw[start:].tolist()}
62
+
63
+
64
+ def _encode_image(
65
+ pipe: Any, image: Any, width: int, height: int, device: str, generator: Any
66
+ ) -> Any:
67
+ """The input image as packed, normalized Krea 2 latents. The Qwen-Image VAE is a video codec, so
68
+ the pixels carry a length-1 temporal axis. The VAE sample is drawn from the run's seeded
69
+ generator, not the global RNG, so the same seed reproduces the same img2img result."""
70
+ import numpy as np
71
+
72
+ resized = image.convert("RGB").resize((width, height))
73
+ array = np.asarray(resized, dtype="float32") / 127.5 - 1.0
74
+ dtype = pipe.vae.dtype
75
+ pixels = torch.from_numpy(array).permute(2, 0, 1)[None, :, None].to(device, dtype)
76
+
77
+ with torch.no_grad():
78
+ latents = pipe.vae.encode(pixels).latent_dist.sample(generator=generator)
79
+ mean, std = _latent_stats(pipe.vae, latents)
80
+ latents = ((latents - mean) / std).squeeze(2) # drop the temporal axis
81
+ batch, channels, latent_h, latent_w = latents.shape
82
+ return pipe._pack_latents(latents, batch, channels, latent_h, latent_w)
83
+
84
+
85
+ def _latent_stats(vae: Any, latents: Any) -> tuple[Any, Any]:
86
+ shape = (1, vae.config.z_dim, 1, 1, 1)
87
+ mean = torch.tensor(vae.config.latents_mean, device=latents.device, dtype=latents.dtype)
88
+ std = torch.tensor(vae.config.latents_std, device=latents.device, dtype=latents.dtype)
89
+ return mean.view(shape), std.view(shape)
@@ -0,0 +1,57 @@
1
+ """Krea 2's answer to "what do I need on disk" - one provider per node variant."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from pathlib import Path
6
+ from typing import Any
7
+
8
+ from ...config import models_dir
9
+ from ..requirements import ModelComponent
10
+ from .requirements import (
11
+ footprint_bytes,
12
+ krea2_requirements,
13
+ resolve_diffusion,
14
+ resolve_text_encoder,
15
+ resolve_vae,
16
+ )
17
+
18
+
19
+ class Krea2Provider:
20
+ """Requirements + fit estimate for one Krea 2 node (``turbo`` or ``raw``)."""
21
+
22
+ def __init__(self, variant: str) -> None:
23
+ self._variant = variant
24
+
25
+ def components(self, params: dict[str, object] | None = None) -> list[ModelComponent]:
26
+ return krea2_requirements(self._variant, params)
27
+
28
+ def download_target(self, component: ModelComponent) -> Path:
29
+ return models_dir() / component.category
30
+
31
+ def estimate(self, policy: Any) -> dict[str, Any] | None:
32
+ """Whether the model will fit this machine, so the popup can warn before a 26GB load.
33
+ ``None`` whenever it can't be sized - a wrong estimate is worse than none."""
34
+ if policy is None:
35
+ return None
36
+ try:
37
+ from ...device.policy import ModelFootprint
38
+ except ImportError:
39
+ return None
40
+ footprint = ModelFootprint(
41
+ **footprint_bytes(
42
+ resolve_diffusion(self._variant), resolve_vae(None), resolve_text_encoder(None)
43
+ )
44
+ )
45
+ fit = policy.estimate_fit(footprint) # pure - never mutates the shared policy
46
+ if fit is None:
47
+ return None
48
+ soft = not fit.fits or fit.plan in ("int8", "offload")
49
+ return {
50
+ "plan": fit.plan,
51
+ "fits": fit.fits,
52
+ "requiredVramMb": int(fit.required_vram_gb * 1024),
53
+ "totalVramMb": int(fit.total_vram_gb * 1024) if fit.total_vram_gb else None,
54
+ "freeVramMb": policy.free_vram_mb(),
55
+ "freeRamMb": policy.free_ram_mb(),
56
+ "warning": fit.note if soft else None,
57
+ }