Krea2: VAE в edit-энкодере и латент под вход

Подключил VAE к TextEncodeQwenImageEdit, добавил латент под пропорции входного изображения. Поставил 10 шагов dpmpp_2m/karras.

README.md

@@ -285,8 +285,9 @@ docker compose up -d
 
 - `workflows/txt2img.json` — Z-Image Turbo: `EmptySD3LatentImage → KSampler → VAEDecode → SaveImage`.
 - `workflows/img2img.json` — Krea2 edit: `LoadImage → TextEncodeQwenImageEdit → KSamplerAdvanced → VAEDecode → SaveImage`.
-- `workflows/mapping.json` — ID узлов, в которые подставляются prompt, negative, seed и входное изображение.
+- `workflows/mapping.json` — ID узлов, в которые подставляются prompt, negative, seed, входное изображение и размер латента.
 - `workflows/Krea2.json` и `workflows/ZImage.json` — исходные экспорты.
+- Krea2 edit: VAE подключён к `TextEncodeQwenImageEdit`, латент подгоняется под пропорции входного изображения; 10 шагов, `dpmpp_2m`/`karras`.
 
 Если меняются workflow, обновите `mapping.json`. При старте бот проверяет, что все узлы и поля существуют.
 

queue_manager.py

@@ -10,15 +10,23 @@ from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from pathlib import Path
 
+from PIL import Image, ImageOps
+
 from comfy_client import ComfyClient, ComfyError
 from config import Settings
 from db import Database, Request, UserAlreadyQueued
 from moderation import PromptBlocked, is_prompt_allowed
 from prompt_enhancer import PromptEnhancer
-from workflows import WorkflowBuilder
+from workflows import WorkflowBuilder, fit_latent_size
 
 logger = logging.getLogger(__name__)
 
+
+def _image_size(path: str) -> tuple[int, int]:
+    with Image.open(path) as image:
+        oriented = ImageOps.exif_transpose(image)
+        return fit_latent_size(*oriented.size)
+
 EventCallback = Callable[["QueueEvent"], Awaitable[None] | None]
 
 
@@ -362,11 +370,16 @@ class QueueManager:
         suffix = input_path.suffix.lower() or ".png"
         remote_name = f"{request.id}_input{suffix}"
         image_name = await self._comfy.upload_image(input_path, remote_name)
+        width, height = await asyncio.to_thread(
+            _image_size, request.input_file_path
+        )
         return self._workflows.build_img2img(
             prompt=prompt,
             negative=negative,
             seed=seed,
             image_name=image_name,
+            width=width,
+            height=height,
         )
 
     async def _refresh(self, request_id: int) -> Request:

tests/test_queue.py

@@ -6,6 +6,7 @@ import time
 from pathlib import Path
 
 import pytest
+from PIL import Image
 
 from config import Settings
 from db import Database, UserAlreadyQueued
@@ -183,7 +184,7 @@ async def test_image_to_image_uploads_input(db: Database, tmp_path: Path) -> Non
     comfy = queue._comfy  # type: ignore[attr-defined]
 
     input_path = tmp_path / "input.png"
-    input_path.write_bytes(b"input")
+    Image.new("RGB", (64, 64), (120, 40, 40)).save(input_path)
 
     await queue.start()
     try:

tests/test_workflows.py

@@ -5,7 +5,7 @@ from pathlib import Path
 
 import pytest
 
-from workflows import WorkflowBuilder, WorkflowError
+from workflows import WorkflowBuilder, WorkflowError, fit_latent_size
 
 BASE = Path(__file__).resolve().parents[1] / "workflows"
 
@@ -39,10 +39,50 @@ def test_build_img2img() -> None:
     )
 
     assert workflow["25"]["inputs"]["prompt"] == "make blue"
+    assert workflow["25"]["inputs"]["vae"] == ["19", 2]
     assert workflow["10"]["inputs"]["text"] == "bad"
     assert workflow["3"]["inputs"]["noise_seed"] == 7
+    assert workflow["3"]["inputs"]["steps"] == 10
+    assert workflow["3"]["inputs"]["sampler_name"] == "dpmpp_2m"
+    assert workflow["3"]["inputs"]["scheduler"] == "karras"
     assert workflow["20"]["class_type"] == "LoadImage"
     assert workflow["20"]["inputs"]["image"] == "input.png"
+    assert workflow["15"]["inputs"]["width"] == 1024
+    assert workflow["15"]["inputs"]["height"] == 1024
+
+
+def test_build_img2img_custom_size() -> None:
+    builder = make_builder()
+    workflow = builder.build_img2img(
+        prompt="make blue",
+        negative="bad",
+        seed=7,
+        image_name="input.png",
+        width=832,
+        height=1216,
+    )
+
+    assert workflow["15"]["inputs"]["width"] == 832
+    assert workflow["15"]["inputs"]["height"] == 1216
+
+
+def test_fit_latent_size_preserves_aspect_ratio() -> None:
+    assert fit_latent_size(1024, 1024) == (1024, 1024)
+
+    width, height = fit_latent_size(1080, 1920)
+    assert width % 16 == 0
+    assert height % 16 == 0
+    assert height > width
+    assert max(width, height) <= 1536
+
+    width, height = fit_latent_size(4000, 3000)
+    assert width % 16 == 0
+    assert height % 16 == 0
+    assert width > height
+
+
+def test_fit_latent_size_fallback_for_invalid_size() -> None:
+    assert fit_latent_size(0, 100) == (1024, 1024)
 
 
 def test_templates_are_not_mutated() -> None:

workflows.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import copy
 import json
+import math
 from collections.abc import Callable
 from pathlib import Path
 from typing import Any
@@ -21,6 +22,7 @@ ROLE_FIELDS: dict[str, dict[str, RoleField]] = {
         "negative": ("text", lambda value: isinstance(value, str)),
         "seed": ("noise_seed", lambda value: isinstance(value, int)),
         "input_image": ("image", lambda value: isinstance(value, str)),
+        "latent": (None, None),
         "save": (None, None),
     },
 }
@@ -28,6 +30,31 @@ ROLE_FIELDS: dict[str, dict[str, RoleField]] = {
 LATENT_CLASSES = {"EmptySD3LatentImage", "EmptyLatentImage", "EmptyLatentImagePresets"}
 
 
+def fit_latent_size(
+    width: int,
+    height: int,
+    *,
+    target_pixels: int = 1024 * 1024,
+    multiple: int = 16,
+    max_side: int = 1536,
+) -> tuple[int, int]:
+    if width <= 0 or height <= 0:
+        return 1024, 1024
+
+    ratio = math.sqrt(target_pixels / (width * height))
+    new_width = width * ratio
+    new_height = height * ratio
+    longest = max(new_width, new_height)
+    if longest > max_side:
+        scale = max_side / longest
+        new_width *= scale
+        new_height *= scale
+
+    new_width = max(multiple, round(new_width / multiple) * multiple)
+    new_height = max(multiple, round(new_height / multiple) * multiple)
+    return int(new_width), int(new_height)
+
+
 class WorkflowError(RuntimeError):
     pass
 
@@ -109,6 +136,8 @@ class WorkflowBuilder:
         negative: str,
         seed: int,
         image_name: str,
+        width: int = 1024,
+        height: int = 1024,
     ) -> dict[str, Any]:
         workflow = copy.deepcopy(self._templates["img2img"])
         mapping = self._mapping["img2img"]
@@ -116,4 +145,7 @@ class WorkflowBuilder:
         workflow[mapping["negative"]]["inputs"]["text"] = negative
         workflow[mapping["seed"]]["inputs"]["noise_seed"] = seed
         workflow[mapping["input_image"]]["inputs"]["image"] = image_name
+        latent = workflow[mapping["latent"]]["inputs"]
+        latent["width"] = width
+        latent["height"] = height
         return workflow

workflows/Krea2.json

@@ -3,10 +3,10 @@
     "inputs": {
       "add_noise": "enable",
       "noise_seed": 782408632091183,
-      "steps": 8,
+      "steps": 10,
       "cfg": 1,
-      "sampler_name": "euler",
+      "sampler_name": "dpmpp_2m",
-      "scheduler": "simple",
+      "scheduler": "karras",
       "start_at_step": 0,
       "end_at_step": 10000,
       "return_with_leftover_noise": "disable",
@@ -116,6 +116,10 @@
         "19",
         1
       ],
+      "vae": [
+        "19",
+        2
+      ],
       "image": [
         "20",
         0

workflows/img2img.json

@@ -3,10 +3,10 @@
     "inputs": {
       "add_noise": "enable",
       "noise_seed": 0,
-      "steps": 8,
+      "steps": 10,
       "cfg": 1,
-      "sampler_name": "euler",
+      "sampler_name": "dpmpp_2m",
-      "scheduler": "simple",
+      "scheduler": "karras",
       "start_at_step": 0,
       "end_at_step": 10000,
       "return_with_leftover_noise": "disable",
@@ -115,6 +115,10 @@
         "19",
         1
       ],
+      "vae": [
+        "19",
+        2
+      ],
       "image": [
         "20",
         0

workflows/mapping.json

@@ -11,6 +11,7 @@
     "negative": "10",
     "seed": "3",
     "input_image": "20",
+    "latent": "15",
     "save": "24"
   }
 }