Update app.py

app.py

@@ -40,6 +40,101 @@ def get_unconditional_conditioning(N, obj_thr, device):
     uc = single_uc.unsqueeze(-1).repeat(1, 1, 1, obj_thr)
     return {"pch_code": uc}
 
+def process_pairs_multiple(mask, tar_image, patch_dir, counter=0, max_ratio=0.8):
+    # 1. Process Reference Object (View)
+    view = cv2.imread(patch_dir)
+    view = cv2.cvtColor(view, cv2.COLOR_BGR2RGB)
+    view = pad_to_square(view, pad_value=255, random=False)
+    view = cv2.resize(view.astype(np.uint8), (224, 224))
+    view = view.astype(np.float32) / 255.0
+
+    # 2. BBox and Mask Logic
+    box_yyxx = get_bbox_from_mask(mask)
+    
+    # Define crop area (using full image here)
+    H1, W1 = tar_image.shape[0], tar_image.shape[1]
+    box_yyxx_crop = [0, H1, 0, W1]
+    
+    # Handle box within crop
+    y1, y2, x1, x2 = box_in_box(box_yyxx, box_yyxx_crop)
+
+    # 3. Create Collage (Input Hint)
+    # Background with hole (zeroed out at object position)
+    collage = tar_image.copy()
+    source_collage = collage.copy()
+    collage[y1:y2, x1:x2, :] = 0
+
+    # Binary mask for the current object hole
+    collage_mask = np.zeros_like(tar_image, dtype=np.float32)
+    collage_mask[y1:y2, x1:x2, :] = 1.0
+
+    # 4. Square Padding & Resizing
+    # Pad all to square (pad_value 2 for mask indicates padding area)
+    tar_square = pad_to_square(tar_image, pad_value=0, random=False)
+    collage_square = pad_to_square(collage, pad_value=0, random=False)
+    mask_square = pad_to_square(collage_mask, pad_value=2, random=False)
+    
+    H2, W2 = collage_square.shape[0], collage_square.shape[1]
+
+    # Resize to model input size
+    tar_res = cv2.resize(tar_square, (512, 512)).astype(np.float32)
+    col_res = cv2.resize(collage_square, (512, 512)).astype(np.float32)
+    mask_res = cv2.resize(mask_square, (512, 512), interpolation=cv2.INTER_NEAREST).astype(np.float32)
+
+    # 5. Mask Value Normalization
+    # Original logic: mask=1 for object, 0 for background, -1 for padding
+    mask_res[mask_res == 2] = -1
+    
+    # For conditioning: keep a 0/1 version for cross-attn mask
+    c_mask = np.where(mask_res[..., 0:1] == 1, 1.0, 0.0).astype(np.float32)
+
+    # 6. Final Item Assembly
+    # Normalize images to [-1, 1]
+    tar_res = tar_res / 127.5 - 1.0
+    col_res = col_res / 127.5 - 1.0
+    
+    # Hint: Concatenate background with the (-1, 0, 1) mask
+    hint_final = np.concatenate([col_res, mask_res[..., :1]], axis=-1)
+
+    item = {
+        f'view{counter}': view,
+        f'hint{counter}': hint_final,
+        f'mask{counter}': c_mask,
+        f'hint_sizes{counter}': np.array([y1, x1, y2, x2]),
+        'jpg': tar_res, # Targets are same for all counters in a pair
+        'collage': source_collage,
+        'extra_sizes': np.array([H1, W1, H2, W2])
+    }
+
+    return item
+
+
+def process_composition(item, obj_thr):
+    collage = item['collage'].copy()
+    collage_mask = np.zeros((collage.shape[0], collage.shape[1], 1), dtype=np.float32)
+
+    for i in reversed(range(obj_thr)):
+        y1, x1, y2, x2 = item['hint_sizes'+str(i)]
+        collage[y1:y2, x1:x2, :] = 0
+        collage_mask[y1:y2,x1:x2,:] = 1.0
+
+    collage = pad_to_square(collage, pad_value = 0, random = False).astype(np.uint8)
+    
+    collage_mask = pad_to_square(collage_mask, pad_value = 2, random = False).astype(np.float32)
+
+    collage = cv2.resize(collage.astype(np.uint8), (512, 512)).astype(np.float32) / 127.5 - 1.0 
+    collage_mask = cv2.resize(collage_mask, (512, 512), interpolation=cv2.INTER_NEAREST).astype(np.float32)
+    
+    if len(collage_mask.shape) == 2: 
+        collage_mask = collage_mask[..., None]
+    
+    collage_mask[collage_mask == 2] = -1.0
+
+    collage_final = np.concatenate([collage, collage_mask[:,:,:1]] , -1)
+    
+    item.update({'hint': collage_final.copy()})
+    return item
+
 # 4. 核心推理函数 (完全照抄你的逻辑)
 @spaces.GPU(duration=120)
 def pics_pairwise_inference(background, img_a, mask_a, img_b, mask_b):