From ff935165a8b037a21a4faa56db2bfcaffc0454c5 Mon Sep 17 00:00:00 2001 From: logtd Date: Fri, 23 Feb 2024 13:33:39 -0700 Subject: [PATCH] fix sub_idxs --- conditioning/embeddings.py | 294 +++++++++--------- ...instance_diffusion_tracking_prompt_node.py | 124 ++++---- 2 files changed, 205 insertions(+), 213 deletions(-) diff --git a/conditioning/embeddings.py b/conditioning/embeddings.py index 6b120fa..ca3f4fd 100644 --- a/conditioning/embeddings.py +++ b/conditioning/embeddings.py @@ -10,194 +10,190 @@ N_MAX_OBJECTS = 30 def get_point_from_box(bbox): - x0, y0, x1, y1 = bbox[0], bbox[1], bbox[2], bbox[3] - return [(x0 + x1) / 2.0, (y0 + y1) / 2.0] + x0, y0, x1, y1 = bbox[0], bbox[1], bbox[2], bbox[3] + return [(x0 + x1) / 2.0, (y0 + y1) / 2.0] def get_empty_binary_mask(img_width, img_height): - return np.zeros((img_width, img_height, 1)) + return np.zeros((img_width, img_height, 1)) def sample_random_points_from_mask(mask, k=N_SCRIBBLE_POINTS): - mask = mask[:, :, 0] - # Find the coordinates of non-zero pixels in the binary mask - nonzero_coords = np.transpose(np.nonzero(mask)) + mask = mask[:, :, 0] + # Find the coordinates of non-zero pixels in the binary mask + nonzero_coords = np.transpose(np.nonzero(mask)) - # Randomly sample 'k' points - # return all zeros if there is no non-zero pixel - if len(nonzero_coords) == 0: - xy_points = [0 for _ in range(k * 2)] + # Randomly sample 'k' points + # return all zeros if there is no non-zero pixel + if len(nonzero_coords) == 0: + xy_points = [0 for _ in range(k * 2)] + return xy_points + + # randomly sample with replacement if there are not enough non-zero pixels + if len(nonzero_coords) < k and len(nonzero_coords) > 0: + random_indices = np.random.choice(len(nonzero_coords), k, replace=True) + # randomly sample withiout replacement if there are enough non-zero pixels + else: + random_indices = np.random.choice( + len(nonzero_coords), k, replace=False) + sampled_points = nonzero_coords[random_indices] + + # order the points by their distance to (0, 0) + # center = np.array([mask.shape[0] // 2, mask.shape[1] // 2]) + center = np.array([0, 0]) + sampled_points = sorted(sampled_points, key=lambda x: np.linalg.norm( + np.array(x) - center)) # np.linalg.norm + + # concatenate x and y coordinates and return them as a list + # [x1,y1,x2,y2,...,x_k,y_k] + xy_points = [] + for x in sampled_points: + xy_points.append(float(x[1])) + xy_points.append(float(x[0])) return xy_points - # randomly sample with replacement if there are not enough non-zero pixels - if len(nonzero_coords) < k and len(nonzero_coords) > 0: - random_indices = np.random.choice(len(nonzero_coords), k, replace=True) - # randomly sample withiout replacement if there are enough non-zero pixels - else: - random_indices = np.random.choice( - len(nonzero_coords), k, replace=False) - sampled_points = nonzero_coords[random_indices] - - # order the points by their distance to (0, 0) - # center = np.array([mask.shape[0] // 2, mask.shape[1] // 2]) - center = np.array([0, 0]) - sampled_points = sorted(sampled_points, key=lambda x: np.linalg.norm( - np.array(x) - center)) # np.linalg.norm - - # concatenate x and y coordinates and return them as a list - # [x1,y1,x2,y2,...,x_k,y_k] - xy_points = [] - for x in sampled_points: - xy_points.append(float(x[1])) - xy_points.append(float(x[0])) - return xy_points - def convert_points(points, img_width, img_height): - # convert polygons/scribbless' coordinates to the relative values (0, 1) - for i in range(len(points)): - if i % 2 == 0: - points[i] = min(points[i] / img_width, 1.0) - else: - points[i] = min(points[i] / img_height, 1.0) - return points + # convert polygons/scribbless' coordinates to the relative values (0, 1) + for i in range(len(points)): + if i % 2 == 0: + points[i] = min(points[i] / img_width, 1.0) + else: + points[i] = min(points[i] / img_height, 1.0) + return points def sample_sparse_points_from_mask(mask, k=256): - n_points = k - n_polygons = n_points // 2 # half points should be sampled from the polygons - mask = mask[:, :, 0] - # sample sparse points from the polygons (boundary) - polygons = binary_mask_to_polygon(mask, tolerance=0.0) - # concatenate polygons to a single list - polygons_single = [] - for polygon in polygons: - polygons_single += polygon - if len(polygons_single) != 0: - # uniformly sample points from the polygon - polygons_single = np.array(polygons_single).reshape(-1, 2) - indexes = np.linspace(0, polygons_single.shape[0] - 1, n_polygons) - indexes = list([int(i) for i in indexes]) + n_points = k + n_polygons = n_points // 2 # half points should be sampled from the polygons + mask = mask[:, :, 0] + # sample sparse points from the polygons (boundary) + polygons = binary_mask_to_polygon(mask, tolerance=0.0) + # concatenate polygons to a single list + polygons_single = [] + for polygon in polygons: + polygons_single += polygon + if len(polygons_single) != 0: + # uniformly sample points from the polygon + polygons_single = np.array(polygons_single).reshape(-1, 2) + indexes = np.linspace(0, polygons_single.shape[0] - 1, n_polygons) + indexes = list([int(i) for i in indexes]) - polygons_single = polygons_single[indexes] - sampled_polygons = [(x[0], x[1]) for x in polygons_single] - else: - return [0 for _ in range(256 * 2)] + polygons_single = polygons_single[indexes] + sampled_polygons = [(x[0], x[1]) for x in polygons_single] + else: + return [0 for _ in range(256 * 2)] - # sample sparse points from the mask - n_inside_points = n_points - len(sampled_polygons) - inside_points = sample_uniform_sparse_points(mask, n_inside_points) + # sample sparse points from the mask + n_inside_points = n_points - len(sampled_polygons) + inside_points = sample_uniform_sparse_points(mask, n_inside_points) - # combine inside_points and sampled_polygons - xy_points = inside_points + sampled_polygons + # combine inside_points and sampled_polygons + xy_points = inside_points + sampled_polygons - # order the points by their distance to (0, 0) - center = np.array([0, 0]) - xy_points = sorted(xy_points, key=lambda x: np.linalg.norm( - np.array(x) - center)) # np.linalg.norm + # order the points by their distance to (0, 0) + center = np.array([0, 0]) + xy_points = sorted(xy_points, key=lambda x: np.linalg.norm( + np.array(x) - center)) # np.linalg.norm - # return the sampled points - sampled_points = [] - for x in xy_points: - sampled_points.append(x[0]) - sampled_points.append(x[1]) - return sampled_points + # return the sampled points + sampled_points = [] + for x in xy_points: + sampled_points.append(x[0]) + sampled_points.append(x[1]) + return sampled_points # [x0, y0, x1, y1] def get_grounding_input_from_coords(coords, img_width, img_height): - x0, y0, x1, y1, coord_width, coord_height = coords - location = [x0 / coord_width, y0 / coord_height, - x1 / coord_width, y1 / coord_height] + x0, y0, x1, y1, coord_width, coord_height = coords + location = [x0 / coord_width, y0 / coord_height, + x1 / coord_width, y1 / coord_height] - point = get_point_from_box(location) - binary_mask = get_empty_binary_mask(img_width, img_height) + point = get_point_from_box(location) + binary_mask = get_empty_binary_mask(img_width, img_height) - scribble = sample_random_points_from_mask(binary_mask, k=N_SCRIBBLE_POINTS) - scribble = convert_points(scribble, img_width, img_height) + scribble = sample_random_points_from_mask(binary_mask, k=N_SCRIBBLE_POINTS) + scribble = convert_points(scribble, img_width, img_height) - polygon = sample_sparse_points_from_mask(binary_mask, k=N_POLYGON_POINTS) - polygon = convert_points(polygon, img_width, img_height) + polygon = sample_sparse_points_from_mask(binary_mask, k=N_POLYGON_POINTS) + polygon = convert_points(polygon, img_width, img_height) - segment = resize(binary_mask.astype(np.float32), - (img_width, img_height)).squeeze() - # segment = np.stack(segment).astype(np.float32).squeeze() if len(segment) > 0 else segment + segment = resize(binary_mask.astype(np.float32), + (img_width, img_height)).squeeze() + # segment = np.stack(segment).astype(np.float32).squeeze() if len(segment) > 0 else segment - return dict( - polygon=polygon, - scribble=scribble, - segment=segment, - box=location, - point=point, - ) + return dict( + polygon=polygon, + scribble=scribble, + segment=segment, + box=location, + point=point, + ) def create_zero_input_tensors(n_frames, img_width, img_height): - masks = torch.zeros(n_frames, N_MAX_OBJECTS) - text_masks = torch.zeros(n_frames, N_MAX_OBJECTS) - text_embeddings = torch.zeros(n_frames, N_MAX_OBJECTS, 768) - box_embeddings = torch.zeros(n_frames, N_MAX_OBJECTS, 4) - polygon_embeddings = torch.zeros( - n_frames, N_MAX_OBJECTS, N_POLYGON_POINTS * 2) - scribble_embeddings = torch.zeros( - n_frames, N_MAX_OBJECTS, N_SCRIBBLE_POINTS * 2) - segment_embeddings = torch.zeros( - n_frames, N_MAX_OBJECTS, img_width, img_height) # TODO: width height order - point_embeddings = torch.zeros(n_frames, N_MAX_OBJECTS, 2) + masks = torch.zeros(n_frames, N_MAX_OBJECTS) + text_masks = torch.zeros(n_frames, N_MAX_OBJECTS) + text_embeddings = torch.zeros(n_frames, N_MAX_OBJECTS, 768) + box_embeddings = torch.zeros(n_frames, N_MAX_OBJECTS, 4) + polygon_embeddings = torch.zeros( + n_frames, N_MAX_OBJECTS, N_POLYGON_POINTS * 2) + scribble_embeddings = torch.zeros( + n_frames, N_MAX_OBJECTS, N_SCRIBBLE_POINTS * 2) + segment_embeddings = torch.zeros( + n_frames, N_MAX_OBJECTS, img_width, img_height) # TODO: width height order + point_embeddings = torch.zeros(n_frames, N_MAX_OBJECTS, 2) - return dict( - masks=masks, - text_masks=text_masks, - prompts=text_embeddings, - boxes=box_embeddings, - polygons=polygon_embeddings, - scribbles=scribble_embeddings, - segments=segment_embeddings, - points=point_embeddings - ) + return dict( + masks=masks, + text_masks=text_masks, + prompts=text_embeddings, + boxes=box_embeddings, + polygons=polygon_embeddings, + scribbles=scribble_embeddings, + segments=segment_embeddings, + points=point_embeddings + ) def get_attn_mask(img_size=64): - return torch.zeros(N_MAX_OBJECTS, img_size, img_size) + return torch.zeros(N_MAX_OBJECTS, img_size, img_size) def prepare_embeddings(conds, latent_shape, idxs, use_masked_att=False): - batch_size, _, latent_height, latent_width = latent_shape - embeddings = create_zero_input_tensors( - batch_size, latent_width, latent_height) - if use_masked_att: - embeddings['att_masks'] = torch.zeros( - batch_size, N_MAX_OBJECTS, latent_width, latent_height) + batch_size, _, latent_height, latent_width = latent_shape + embeddings = create_zero_input_tensors( + batch_size, latent_width, latent_height) + if use_masked_att: + embeddings['att_masks'] = torch.zeros( + batch_size, N_MAX_OBJECTS, latent_width, latent_height) - for idx in idxs: - for cond_idx, cond in enumerate(conds): - if cond['positions'][idx] is None: - continue + for grounding_idx, frame_idx in enumerate(idxs): + for cond_idx, cond in enumerate(conds): + if cond['positions'][frame_idx] is None: + continue - grounding = get_grounding_input_from_coords( - cond['positions'][idx], latent_width, latent_height) - embeddings['masks'][idx][cond_idx] = 1 - embeddings['text_masks'][idx][cond_idx] = 1 - embeddings['prompts'][idx][cond_idx] = cond['cond_pooled'] - embeddings['boxes'][idx][cond_idx] = torch.tensor( - grounding['box']) - embeddings['polygons'][idx][cond_idx] = torch.tensor( - grounding['polygon']) - embeddings['scribbles'][idx][cond_idx] = torch.tensor( - grounding['scribble']) - embeddings['segments'][idx][cond_idx] = torch.tensor( - grounding['segment']) - embeddings['points'][idx][cond_idx] = torch.tensor( - grounding['point']) + grounding = get_grounding_input_from_coords( + cond['positions'][frame_idx], latent_width, latent_height) + embeddings['masks'][grounding_idx][cond_idx] = 1 + embeddings['text_masks'][grounding_idx][cond_idx] = 1 + embeddings['prompts'][grounding_idx][cond_idx] = cond['cond_pooled'] + embeddings['boxes'][grounding_idx][cond_idx] = torch.tensor( + grounding['box']) + embeddings['polygons'][grounding_idx][cond_idx] = torch.tensor( + grounding['polygon']) + embeddings['scribbles'][grounding_idx][cond_idx] = torch.tensor( + grounding['scribble']) + embeddings['segments'][grounding_idx][cond_idx] = torch.tensor( + grounding['segment']) + embeddings['points'][grounding_idx][cond_idx] = torch.tensor( + grounding['point']) - if use_masked_att: - box = grounding['box'] - x1, y1, x2, y2 = int(np.round(box[0] * latent_width)), int(np.round(box[1] * latent_height)), int( - np.round(box[2] * latent_width)), int(np.round(box[3] * latent_height)) - embeddings['att_masks'][idx][cond_idx][x1:x2, y1:y2] = 1 + if use_masked_att: + box = grounding['box'] + x1, y1, x2, y2 = int(np.round(box[0] * latent_width)), int(np.round(box[1] * latent_height)), int( + np.round(box[2] * latent_width)), int(np.round(box[3] * latent_height)) + embeddings['att_masks'][grounding_idx][cond_idx][x1:x2, y1:y2] = 1 - return embeddings - - -def get_model_inputs(): - return + return embeddings diff --git a/nodes/instance_diffusion_tracking_prompt_node.py b/nodes/instance_diffusion_tracking_prompt_node.py index e56d3e0..f30a9d9 100644 --- a/nodes/instance_diffusion_tracking_prompt_node.py +++ b/nodes/instance_diffusion_tracking_prompt_node.py @@ -3,79 +3,75 @@ from ..conditioning.instance_conditioning import InstanceConditioning class InstanceDiffusionTrackingPromptNode: - @classmethod - def INPUT_TYPES(s): - return {"required": {"positive": ("CONDITIONING", ), - "negative": ("CONDITIONING", ), - "clip": ("CLIP", ), - "tracking": ("TRACKING", ), - "positionnet": ("POSITIONNET", ), - "fusers": ("FUSERS", ), - "fusers_batch_size": ("INT", {"default": 5, "min": 0, "max": 10000, "step": 8}), - "positive_text": ("STRING", {"multiline": True}), - "negative_text": ("STRING", {"multiline": True}), - }} - RETURN_TYPES = ("CONDITIONING", "CONDITIONING") - RETURN_NAMES = ("positive", "negative") - FUNCTION = "append" + @classmethod + def INPUT_TYPES(s): + return {"required": {"positive": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + "clip": ("CLIP", ), + "tracking": ("TRACKING", ), + "positionnet": ("POSITIONNET", ), + "fusers": ("FUSERS", ), + "fusers_batch_size": ("INT", {"default": 5, "min": 0, "max": 10000, "step": 8}), + "positive_text": ("STRING", {"multiline": True}), + "negative_text": ("STRING", {"multiline": True}), + }} + RETURN_TYPES = ("CONDITIONING", "CONDITIONING") + RETURN_NAMES = ("positive", "negative") + FUNCTION = "append" - CATEGORY = "instance/conditioning" + CATEGORY = "instance/conditioning" - def _get_position_conds(self, clip, tracking, text): - # Get prompts and their class id and trakcer id - prompt_pairs = extract_prompts(text) + def _get_position_conds(self, clip, tracking, text): + # Get prompts and their class id and trakcer id + prompt_pairs = extract_prompts(text) - # Go through prompt pairs, encode prompts, and join with positions from tracking - position_conds = [] - for tracker_id, class_id, prompt in prompt_pairs: - _, cond_pooled = clip.encode_from_tokens( - clip.tokenize(prompt), return_pooled=True) - # A tracker_id of -1 means that it is prompting all instances of a single class - if tracker_id != -1: - position_cond = {'cond_pooled': cond_pooled, 'positions': - tracking[class_id][tracker_id]} - position_conds.append(position_cond) - else: - for tracker_id in tracking[class_id]: - position_cond = {'cond_pooled': cond_pooled, - 'positions': tracking[class_id][tracker_id]} - position_conds.append(position_cond) + # Go through prompt pairs, encode prompts, and join with positions from tracking + position_conds = [] + for tracker_id, class_id, prompt in prompt_pairs: + _, cond_pooled = clip.encode_from_tokens( + clip.tokenize(prompt), return_pooled=True) + # A tracker_id of -1 means that it is prompting all instances of a single class + if tracker_id != -1: + position_cond = {'cond_pooled': cond_pooled, 'positions': + tracking[class_id][tracker_id]} + position_conds.append(position_cond) + else: + for tracker_id in tracking[class_id]: + position_cond = {'cond_pooled': cond_pooled, + 'positions': tracking[class_id][tracker_id]} + position_conds.append(position_cond) - return position_conds + return position_conds - def _apply_position_conds(self, position_conds, conditioning, fusers, positionnet, fusers_batch_size): - # Add prompts+embeddings to the input conditionings - cond_out = [] - for t in conditioning: - n = [t[0], t[1].copy()] - cond = n[1] - prev = [] - has_instance = 'instance_diffusion' in cond - instance_conditioning = conditioning['instance_diffusion'] if has_instance else InstanceConditioning( - fusers, positionnet, fusers_batch_size) - cond['instance_diffusion'] = instance_conditioning - instance_conditioning.add_conds(position_conds) + def _apply_position_conds(self, position_conds, conditioning, fusers, positionnet, fusers_batch_size): + # Add prompts+embeddings to the input conditionings + cond_out = [] + for t in conditioning: + n = [t[0], t[1].copy()] + cond = n[1] + prev = [] + has_instance = 'instance_diffusion' in cond + instance_conditioning = conditioning['instance_diffusion'] if has_instance else InstanceConditioning( + fusers, positionnet, fusers_batch_size) + cond['instance_diffusion'] = instance_conditioning + instance_conditioning.add_conds(position_conds) - cond['instance_diffusion'] = { - 'conds': prev + position_conds, - } + cond['gligen'] = ('position', instance_conditioning, None) - cond['gligen'] = ('position', instance_conditioning, None) + cond_out.append(n) - cond_out.append(n) + return cond_out - return cond_out + def append(self, positive, negative, clip, tracking, fusers, positionnet, fusers_batch_size, positive_text, negative_text): - def append(self, positive, negative, clip, tracking, fusers, positionnet, fusers_batch_size, positive_text, negative_text): + positive_positions = self._get_position_conds( + clip, tracking, positive_text) + positive = self._apply_position_conds( + positive_positions, positive, fusers, positionnet, fusers_batch_size) - positive_positions = self._get_position_conds( - clip, tracking, positive_text) - positive = self._apply_position_conds( - positive_positions, positive, fusers, positionnet, fusers_batch_size) + negative_positions = self._get_position_conds( + clip, tracking, negative_text) + negative = self._apply_position_conds( + negative_positions, negative, fusers, positionnet, fusers_batch_size) - negative_positions = self._get_position_conds( - clip, tracking, negative_text) - negative = self._apply_position_conds( - negative_positions, negative, fusers, positionnet, fusers_batch_size) - - return (positive, negative) + return (positive, negative)