flux3: multi-reference sampling, match_target_res reference sizing, UI sections

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)

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:
+        if ctrl_imgs:
+            ref_latents = []
+            for ctrl_img in ctrl_imgs:
                 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
+                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),

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",
     ],
   },
   {