Flux Kontext and more

wan/multitalk/attention.py

@@ -333,7 +333,7 @@ class SingleStreamMutiAttention(SingleStreamAttention):
 
         human1 = normalize_and_scale(x_ref_attn_map[0], (human1_min_value, human1_max_value), (self.rope_h1[0], self.rope_h1[1]))
         human2 = normalize_and_scale(x_ref_attn_map[1], (human2_min_value, human2_max_value), (self.rope_h2[0], self.rope_h2[1]))
-        back   = torch.full((x_ref_attn_map.size(1),), self.rope_bak, dtype=human1.dtype).to(human1.device)
+        back   = torch.full((x_ref_attn_map.size(1),), self.rope_bak, dtype=human1.dtype, device=human1.device)
         max_indices = x_ref_attn_map.argmax(dim=0)
         normalized_map = torch.stack([human1, human2, back], dim=1)
         normalized_pos = normalized_map[range(x_ref_attn_map.size(1)), max_indices] # N 
@@ -351,7 +351,7 @@ class SingleStreamMutiAttention(SingleStreamAttention):
         if self.qk_norm:
             encoder_k = self.add_k_norm(encoder_k)
 
-        per_frame = torch.zeros(N_a, dtype=encoder_k.dtype).to(encoder_k.device)
+        per_frame = torch.zeros(N_a, dtype=encoder_k.dtype, device=encoder_k.device)
         per_frame[:per_frame.size(0)//2] = (self.rope_h1[0] + self.rope_h1[1]) / 2
         per_frame[per_frame.size(0)//2:] = (self.rope_h2[0] + self.rope_h2[1]) / 2
         encoder_pos = torch.concat([per_frame]*N_t, dim=0)

wan/multitalk/multitalk.py

@@ -272,6 +272,34 @@ def timestep_transform(
     new_t = new_t * num_timesteps
     return new_t
 
+def parse_speakers_locations(speakers_locations):
+    bbox = {}
+    if speakers_locations is None or len(speakers_locations) == 0:
+        return None, ""
+    speakers = speakers_locations.split(" ")
+    if len(speakers) !=2:
+        error= "Two speakers locations should be defined"
+        return "", error
+    
+    for i, speaker in enumerate(speakers):
+        location = speaker.strip().split(":")
+        if len(location) not in (2,4):
+            error = f"Invalid Speaker Location '{location}'. A Speaker Location should be defined in the format Left:Right or usuing a BBox Left:Top:Right:Bottom"
+            return "", error
+        try:
+            good = False
+            location_float = [ float(val) for val in location]
+            good = all( 0 <= val <= 100 for val in location_float)
+        except:
+            pass
+        if not good:
+            error = f"Invalid Speaker Location '{location}'. Each number should be between 0 and 100."
+            return "", error
+        if len(location_float) == 2:
+            location_float = [location_float[0], 0, location_float[1], 100]
+        bbox[f"human{i}"] = location_float
+    return bbox, ""
+
 
 # construct human mask
 def get_target_masks(HUMAN_NUMBER, lat_h, lat_w, src_h, src_w, face_scale = 0.05, bbox = None):
@@ -286,7 +314,9 @@ def get_target_masks(HUMAN_NUMBER, lat_h, lat_w, src_h, src_w, face_scale = 0.05
             assert len(bbox) == HUMAN_NUMBER, f"The number of target bbox should be the same with cond_audio"
             background_mask = torch.zeros([src_h, src_w])
             for _, person_bbox in bbox.items():
-                x_min, y_min, x_max, y_max = person_bbox
+                y_min, x_min, y_max, x_max = person_bbox
+                x_min, y_min, x_max, y_max = max(x_min,5), max(y_min, 5), min(x_max,95), min(y_max,95)                
+                x_min, y_min, x_max, y_max =  int(src_h * x_min / 100), int(src_w * y_min / 100), int(src_h * x_max / 100), int(src_w * y_max / 100)
                 human_mask = torch.zeros([src_h, src_w])
                 human_mask[int(x_min):int(x_max), int(y_min):int(y_max)] = 1
                 background_mask += human_mask
@@ -306,7 +336,7 @@ def get_target_masks(HUMAN_NUMBER, lat_h, lat_w, src_h, src_w, face_scale = 0.05
             human_masks = [human_mask1, human_mask2]
         background_mask = torch.where(background_mask > 0, torch.tensor(0), torch.tensor(1))
         human_masks.append(background_mask)
-    
+    # toto = Image.fromarray(human_masks[2].mul_(255).unsqueeze(-1).repeat(1,1,3).to(torch.uint8).cpu().numpy())
     ref_target_masks = torch.stack(human_masks, dim=0) #.to(self.device)
     # resize and centercrop for ref_target_masks 
     # ref_target_masks = resize_and_centercrop(ref_target_masks, (target_h, target_w))

wan/multitalk/multitalk_utils.py

@@ -128,7 +128,7 @@ def get_attn_map_with_target(visual_q, ref_k, shape, ref_target_masks=None, spli
 
     _, seq_lens, heads, _ = visual_q.shape
     class_num, _ = ref_target_masks.shape
-    x_ref_attn_maps = torch.zeros(class_num, seq_lens).to(visual_q.device).to(visual_q.dtype)
+    x_ref_attn_maps = torch.zeros(class_num, seq_lens, dtype=visual_q.device, device=visual_q.dtype)
 
     split_chunk = heads // split_num