initial

utils/usdu_managment.py

@@ -36,7 +36,7 @@ JOB_NUM_TILES_PER_IMAGE = 'num_tiles_per_image'  # For static
 TASK_TYPE_TILE = 'tile'
 TASK_TYPE_IMAGE = 'image'
 
-async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_image=None, all_indices=None, enabled_workers=None, task_assignments=None):
+async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_image=None, all_indices=None, enabled_workers=None, task_assignments=None, batched_static: bool = False):
     """Unified initialization for job queues in static and dynamic modes."""
     prompt_server = ensure_tile_jobs_initialized()
     async with prompt_server.distributed_tile_jobs_lock:
@@ -61,12 +61,19 @@ async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_ima
             debug_log(f"Initialized dynamic queue with {batch_size} pending images")
         elif mode == 'static':
             job_data[JOB_NUM_TILES_PER_IMAGE] = num_tiles_per_image
-            # For static with dynamic distribution, populate all tile indices in pending queue
+            job_data[JOB_BATCH_SIZE] = batch_size
+            job_data['batched_static'] = bool(batched_static)
+            # For batched static distribution, populate only tile ids [0..num_tiles_per_image-1]
             pending_queue = job_data[JOB_PENDING_TASKS]
-            total_tiles = batch_size * num_tiles_per_image
-            for i in range(total_tiles):
-                await pending_queue.put(i)
-            debug_log(f"Initialized static queue with {total_tiles} pending tiles for dynamic distribution")
+            if batched_static and num_tiles_per_image is not None:
+                for i in range(num_tiles_per_image):
+                    await pending_queue.put(i)
+                debug_log(f"Initialized static queue (batched) with {num_tiles_per_image} tile ids for dynamic distribution across batch of {batch_size}")
+            else:
+                total_tiles = batch_size * num_tiles_per_image
+                for i in range(total_tiles):
+                    await pending_queue.put(i)
+                debug_log(f"Initialized static queue with {total_tiles} pending tiles for dynamic distribution")
             
             # Keep backward compatibility - if task assignments provided, still track them
             if task_assignments and enabled_workers:
@@ -708,7 +715,8 @@ def clone_control_chain(control, clone_hint=True):
         return None
     new_control = copy.copy(control)  # Shallow copy (shares model)
     if clone_hint and hasattr(control, 'cond_hint_original'):
-        new_control.cond_hint_original = control.cond_hint_original.clone()
+        hint = getattr(control, 'cond_hint_original', None)
+        new_control.cond_hint_original = hint.clone() if hint is not None else None
     if hasattr(control, 'previous_controlnet'):
         new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint)
     return new_control
@@ -722,10 +730,12 @@ def clone_conditioning(cond_list, clone_hints=True):
         if 'control' in new_dict:
             new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints)
         if 'mask' in new_dict:
-            new_dict['mask'] = new_dict['mask'].clone()
+            if new_dict['mask'] is not None:
+                new_dict['mask'] = new_dict['mask'].clone()
         # Handle other potential fields if needed
         if 'pooled_output' in new_dict:
-            new_dict['pooled_output'] = new_dict['pooled_output'].clone()
+            if new_dict['pooled_output'] is not None:
+                new_dict['pooled_output'] = new_dict['pooled_output'].clone()
         if 'area' in new_dict:
             new_dict['area'] = new_dict['area'][:]  # Shallow copy list/tuple
         new_cond.append([new_emb, new_dict])

utils/usdu_utils.py

@@ -4,7 +4,6 @@ import torch
 import torch.nn.functional as F
 from torchvision.transforms import GaussianBlur
 import math
-import os
 
 if (not hasattr(Image, 'Resampling')):  # For older versions of Pillow
     Image.Resampling = Image
@@ -442,6 +441,61 @@ def crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad
 
     cond_dict["mask"] = torch.cat(masks, dim=0)  # (B, H, W)
 
+# Added Flux-Kontext Support crop_reference_latents by TBG ETUR
+def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
+    """
+    1. Resize each latent to `canvas_size` in latent units.
+    2. Crop the rectangle `region` (pixel coordinates).
+    3. Down-sample the crop to latent-space `tile_size`.
+    Expects a list of BCHW tensors under "reference_latents".
+    """
+
+    latents = cond_dict.get("reference_latents")
+    if not isinstance(latents, list):
+        return  # nothing to do
+
+    k = 8  # down-sample factor from pixel space → latent space (SD-type models)
+
+    W_can_px, H_can_px = canvas_size
+    # canvas size expressed in latent units
+    W_can_lat, H_can_lat = W_can_px // k, H_can_px // k
+
+    W_tile_px, H_tile_px = tile_size
+    W_tile_lat, H_tile_lat = max(1, W_tile_px // k), max(1, H_tile_px // k)
+
+    x1_px, y1_px, x2_px, y2_px = region
+
+    new_latents = []
+    for t in latents:  # (B,C,H_lat_in,W_lat_in)
+        if t.ndim != 4:
+            raise ValueError(f"expected BCHW, got {t.shape}")
+
+        # 1. Resize to canvas resolution in latent units only if needed
+        if t.shape[-2:] != (H_can_lat, W_can_lat):
+            t = F.interpolate(t,
+                              size=(H_can_lat, W_can_lat),
+                              mode="bilinear",
+                              align_corners=False)
+
+        # 2. Convert pixel crop → latent slice
+        w0_lat = int(round(x1_px / k))
+        w1_lat = int(round(x2_px / k))
+        h0_lat = int(round(y1_px / k))
+        h1_lat = int(round(y2_px / k))
+
+        cropped = t[:, :, h0_lat:h1_lat, w0_lat:w1_lat]  # view
+
+        # 3. Down-sample to latent-tile size
+        cropped = F.interpolate(cropped,
+                                size=(H_tile_lat, W_tile_lat),
+                                mode="bilinear",
+                                align_corners=False)
+
+        new_latents.append(cropped)
+
+    cond_dict["reference_latents"] = new_latents
+
+
 
 def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0):
     cropped = []
@@ -452,5 +506,6 @@ def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0)
         crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
         crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
         crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
+        crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
         cropped.append(n)
     return cropped