From 34806cf820a557c9561660579e7cb250f212419f Mon Sep 17 00:00:00 2001 From: MpakobecEx Date: Tue, 20 Oct 2026 01:07:07 +0500 Subject: [PATCH] =?UTF-8?q?Krea2:=20VAE=20=D0=B2=20edit-=D1=8D=D0=BD=D0=BA?= =?UTF-8?q?=D0=BE=D0=B4=D0=B5=D1=80=D0=B5=20=D0=B8=20=D0=BB=D0=B0=D1=82?= =?UTF-8?q?=D0=B5=D0=BD=D1=82=20=D0=BF=D0=BE=D0=B4=20=D0=B2=D1=85=D0=BE?= =?UTF-8?q?=D0=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Подключил VAE к TextEncodeQwenImageEdit, добавил латент под пропорции входного изображения. Поставил 10 шагов dpmpp_2m/karras. --- README.md | 3 ++- queue_manager.py | 15 ++++++++++++++- tests/test_queue.py | 3 ++- tests/test_workflows.py | 42 ++++++++++++++++++++++++++++++++++++++++- workflows.py | 32 +++++++++++++++++++++++++++++++ workflows/Krea2.json | 10 +++++++--- workflows/img2img.json | 10 +++++++--- workflows/mapping.json | 1 + 8 files changed, 106 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index f6bd85b..7c46376 100644 --- a/README.md +++ b/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`. При старте бот проверяет, что все узлы и поля существуют. diff --git a/queue_manager.py b/queue_manager.py index 47c94a0..3811c72 100644 --- a/queue_manager.py +++ b/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: diff --git a/tests/test_queue.py b/tests/test_queue.py index dab7c3f..c46c6ad 100644 --- a/tests/test_queue.py +++ b/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: diff --git a/tests/test_workflows.py b/tests/test_workflows.py index ca1fdb5..b4163df 100644 --- a/tests/test_workflows.py +++ b/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: diff --git a/workflows.py b/workflows.py index dd24195..562f398 100644 --- a/workflows.py +++ b/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 diff --git a/workflows/Krea2.json b/workflows/Krea2.json index b8d39bc..3d178ed 100644 --- a/workflows/Krea2.json +++ b/workflows/Krea2.json @@ -3,10 +3,10 @@ "inputs": { "add_noise": "enable", "noise_seed": 782408632091183, - "steps": 8, + "steps": 10, "cfg": 1, - "sampler_name": "euler", - "scheduler": "simple", + "sampler_name": "dpmpp_2m", + "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 diff --git a/workflows/img2img.json b/workflows/img2img.json index 3708c24..1ab4a5f 100644 --- a/workflows/img2img.json +++ b/workflows/img2img.json @@ -3,10 +3,10 @@ "inputs": { "add_noise": "enable", "noise_seed": 0, - "steps": 8, + "steps": 10, "cfg": 1, - "sampler_name": "euler", - "scheduler": "simple", + "sampler_name": "dpmpp_2m", + "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 diff --git a/workflows/mapping.json b/workflows/mapping.json index 61338b6..5528f71 100644 --- a/workflows/mapping.json +++ b/workflows/mapping.json @@ -11,6 +11,7 @@ "negative": "10", "seed": "3", "input_image": "20", + "latent": "15", "save": "24" } }