From dbbb7ee8f4ae9401d53e49d9561b96a85f73d1bd Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 14 Oct 2026 13:37:40 -0600 Subject: [PATCH] flux3: multi-reference sampling, match_target_res reference sizing, UI sections --- .../diffusion_models/flux3/flux3.py | 58 ++++++++++++++++--- .../diffusion_models/flux3/src/pipeline.py | 18 +++--- extensions_built_in/diffusion_models/ui.tsx | 11 +++- 3 files changed, 68 insertions(+), 19 deletions(-) diff --git a/extensions_built_in/diffusion_models/flux3/flux3.py b/extensions_built_in/diffusion_models/flux3/flux3.py index 1a40dc5..3be9205 100644 --- a/extensions_built_in/diffusion_models/flux3/flux3.py +++ b/extensions_built_in/diffusion_models/flux3/flux3.py @@ -25,10 +25,12 @@ target = noise - clean), so no time flips anywhere. """ import json +import math import os from typing import TYPE_CHECKING, List, Optional import torch +import torch.nn.functional as F from PIL import Image from safetensors.torch import load_file, save_file from torchvision.transforms.functional import to_tensor @@ -103,6 +105,11 @@ class Flux3Model(BaseModel): self.model_config.model_kwargs.get("text_pad_multiple", DEFAULT_PAD_MULTIPLE) ) self.package_config = None + # references are area-matched to the target by default; off = keep their + # own size, only scaling down when larger than the target area + self.match_target_res = bool( + self.model_config.model_kwargs.get("match_target_res", True) + ) self.latent_space_version = "flux3_v1" @staticmethod @@ -318,11 +325,23 @@ class Flux3Model(BaseModel): gen_config.negative_prompt or "", target ) - ctrl_img = None - if gen_config.ctrl_img is not None: - img = Image.open(gen_config.ctrl_img).convert("RGB") - img = img.resize((gen_config.width, gen_config.height), Image.LANCZOS) - ctrl_img = to_tensor(img).unsqueeze(0) + # references: ctrl_img_1 mirrors ctrl_img, so dedupe by path order + ctrl_paths = [] + for path in ( + gen_config.ctrl_img, + gen_config.ctrl_img_1, + gen_config.ctrl_img_2, + gen_config.ctrl_img_3, + ): + if path is not None and path not in ctrl_paths: + ctrl_paths.append(path) + target_pixels = gen_config.width * gen_config.height + ctrl_imgs = [ + self._fit_reference( + to_tensor(Image.open(path).convert("RGB")).unsqueeze(0), target_pixels + ) + for path in ctrl_paths + ] alpha = self.model_config.model_kwargs.get("sample_alpha", None) alpha = float(alpha) if alpha is not None else None @@ -336,11 +355,30 @@ class Flux3Model(BaseModel): guidance_scale=gen_config.guidance_scale, latents=gen_config.latents, generator=generator, - ctrl_img=ctrl_img, + ctrl_imgs=ctrl_imgs, alpha=alpha, ) return result[0] + def _fit_reference(self, img: torch.Tensor, target_pixels: int) -> torch.Tensor: + """Size a (1, C, H, W) reference for the DiT. References keep their own + aspect ratio on their own position grid; match_target_res scales the + area to the target's, otherwise the image only scales down when it is + larger than the target area. Snapped to the 16-px grid.""" + sc = self.get_bucket_divisibility() + h, w = img.shape[-2], img.shape[-1] + if self.match_target_res or h * w > target_pixels: + ratio = h / w + new_h = math.sqrt(target_pixels * ratio) + new_w = new_h / ratio + else: + new_h, new_w = float(h), float(w) + new_h = max(sc, int(round(new_h / sc)) * sc) + new_w = max(sc, int(round(new_w / sc)) * sc) + if (new_h, new_w) != (h, w): + img = F.interpolate(img.float(), size=(new_h, new_w), mode="bilinear") + return img + def _encode_at_length(self, prompt: str, length: int) -> AdvancedPromptEmbeds: if self.text_encoder.device == torch.device("cpu"): self.text_encoder.to(self.device_torch) @@ -402,15 +440,17 @@ class Flux3Model(BaseModel): elif batch.control_tensor is not None: control_list = [[batch.control_tensor[b]] for b in range(batch_size)] if control_list is not None: + target_pixels = h_lat * w_lat * 16 * 16 ref_latents = [] for controls in control_list: sample_refs = [] for ctrl in controls: - if ctrl.ndim == 4: - ctrl = ctrl[0] + if ctrl.ndim == 3: + ctrl = ctrl.unsqueeze(0) + ctrl = self._fit_reference(ctrl, target_pixels) # control tensors arrive in [0, 1] lat = self.encode_images( - ctrl.unsqueeze(0) * 2 - 1, device=device, dtype=dtype + ctrl * 2 - 1, device=device, dtype=dtype ) sample_refs.append(lat[0]) ref_latents.append(sample_refs) diff --git a/extensions_built_in/diffusion_models/flux3/src/pipeline.py b/extensions_built_in/diffusion_models/flux3/src/pipeline.py index ab935dc..8001c23 100644 --- a/extensions_built_in/diffusion_models/flux3/src/pipeline.py +++ b/extensions_built_in/diffusion_models/flux3/src/pipeline.py @@ -143,7 +143,7 @@ class Flux3Pipeline: guidance_scale: float = DEFAULT_GUIDANCE, latents: Optional[torch.Tensor] = None, generator: Optional[torch.Generator] = None, - ctrl_img: Optional[torch.Tensor] = None, # (1, 3, H, W) in [0, 1] + ctrl_imgs: Optional[List[torch.Tensor]] = None, # each (1, 3, H, W) in [0, 1] alpha: Optional[float] = None, **kwargs, ): @@ -178,14 +178,16 @@ class Flux3Pipeline: x = x.float() ids = ids.repeat(cfg_batch, 1, 1).to(device) - # reference-image conditioning + # reference-image conditioning: reference i rides at l = -1 - i cond_kwargs = {} - if ctrl_img is not None: - ctrl = ctrl_img.to(device, dtype) * 2 - 1 - cond_latent = model.encode_images(list(ctrl), device=device, dtype=dtype) - xc, xc_ids, xc_ts = pack_i2i_image_condition( - [[cond_latent[0]]], device, dtype - ) + if ctrl_imgs: + ref_latents = [] + for ctrl_img in ctrl_imgs: + ctrl = ctrl_img.to(device, dtype) * 2 - 1 + ref_latents.append( + model.encode_images(list(ctrl), device=device, dtype=dtype)[0] + ) + xc, xc_ids, xc_ts = pack_i2i_image_condition([ref_latents], device, dtype) cond_kwargs = { "x_image_cond": xc.repeat(cfg_batch, 1, 1), "x_image_cond_ids": xc_ids.repeat(cfg_batch, 1, 1), diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index 47b5c89..28afd31 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/extensions_built_in/diffusion_models/ui.tsx @@ -1738,13 +1738,20 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [ "config.process[0].sample.guidance_scale": [3.5, 4], "config.process[0].sample.sample_steps": [28, 25], "config.process[0].train.timestep_type": ["shift", "sigmoid"], + "config.process[0].model.model_kwargs": [ + { + match_target_res: true, + }, + {}, + ], }, disableSections: ["network.conv"], additionalSections: [ - "datasets.control_path", - "sample.ctrl_img", + "datasets.multi_control_paths", + "sample.multi_ctrl_imgs", "model.layer_offloading", "model.low_vram", + "model.qie.match_target_res", ], }, {