cleanup

_utils/attn_utils.py

@@ -1,592 +0,0 @@
-import abc
-
-import cv2
-import numpy as np
-import torch
-from IPython.display import display
-from PIL import Image
-from typing import Union, Tuple, List
-from einops import rearrange, repeat
-import math
-from torch import nn, einsum
-from inspect import isfunction
-from diffusers.utils import logging
-try:
-    from diffusers.models.unet_2d_condition import UNet2DConditionOutput
-except:
-    from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
-
-try:
-    from diffusers.models.cross_attention import CrossAttention
-except:
-    from diffusers.models.attention_processor import Attention as CrossAttention
-
-MAX_NUM_WORDS = 77
-LOW_RESOURCE = False 
-
-class CountingCrossAttnProcessor1:
-
-    def __init__(self, attnstore, place_in_unet):
-        super().__init__()
-        self.attnstore = attnstore
-        self.place_in_unet = place_in_unet
-
-    def __call__(self, attn_layer: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None):
-        batch_size, sequence_length, dim = hidden_states.shape
-        h = attn_layer.heads
-        q = attn_layer.to_q(hidden_states)
-        is_cross = encoder_hidden_states is not None
-        context = encoder_hidden_states if is_cross else hidden_states
-        k = attn_layer.to_k(context)
-        v = attn_layer.to_v(context)
-        # q = attn_layer.reshape_heads_to_batch_dim(q)
-        # k = attn_layer.reshape_heads_to_batch_dim(k)
-        # v = attn_layer.reshape_heads_to_batch_dim(v)
-        # q = attn_layer.head_to_batch_dim(q)
-        # k = attn_layer.head_to_batch_dim(k)
-        # v = attn_layer.head_to_batch_dim(v)
-        q = self.head_to_batch_dim(q, h)
-        k = self.head_to_batch_dim(k, h)
-        v = self.head_to_batch_dim(v, h)
-
-        sim = torch.einsum("b i d, b j d -> b i j", q, k) * attn_layer.scale
-
-        if attention_mask is not None:
-            attention_mask = attention_mask.reshape(batch_size, -1)
-            max_neg_value = -torch.finfo(sim.dtype).max
-            attention_mask = attention_mask[:, None, :].repeat(h, 1, 1)
-            sim.masked_fill_(~attention_mask, max_neg_value)
-
-        # attention, what we cannot get enough of
-        attn_ = sim.softmax(dim=-1).clone()
-        # softmax = nn.Softmax(dim=-1)
-        # attn_ = softmax(sim)
-        self.attnstore(attn_, is_cross, self.place_in_unet)
-        out = torch.einsum("b i j, b j d -> b i d", attn_, v)
-        # out = attn_layer.batch_to_head_dim(out)
-        out = self.batch_to_head_dim(out, h)
-
-        if type(attn_layer.to_out) is torch.nn.modules.container.ModuleList:
-            to_out = attn_layer.to_out[0]
-        else:
-            to_out = attn_layer.to_out
-
-        out = to_out(out)
-        return out
-    
-    def batch_to_head_dim(self, tensor, head_size):
-        # head_size = self.heads
-        batch_size, seq_len, dim = tensor.shape
-        tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)
-        tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)
-        return tensor
-
-    def head_to_batch_dim(self, tensor, head_size, out_dim=3):
-        # head_size = self.heads
-        batch_size, seq_len, dim = tensor.shape
-        tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)
-        tensor = tensor.permute(0, 2, 1, 3)
-
-        if out_dim == 3:
-            tensor = tensor.reshape(batch_size * head_size, seq_len, dim // head_size)
-
-        return tensor
-
-
-def register_attention_control(model, controller):
-
-    attn_procs = {}
-    cross_att_count = 0
-    for name in model.unet.attn_processors.keys():
-        cross_attention_dim = None if name.endswith("attn1.processor") else model.unet.config.cross_attention_dim
-        if name.startswith("mid_block"):
-            hidden_size = model.unet.config.block_out_channels[-1]
-            place_in_unet = "mid"
-        elif name.startswith("up_blocks"):
-            block_id = int(name[len("up_blocks.")])
-            hidden_size = list(reversed(model.unet.config.block_out_channels))[block_id]
-            place_in_unet = "up"
-        elif name.startswith("down_blocks"):
-            block_id = int(name[len("down_blocks.")])
-            hidden_size = model.unet.config.block_out_channels[block_id]
-            place_in_unet = "down"
-        else:
-            continue
-
-        cross_att_count += 1
-        # attn_procs[name] = AttendExciteCrossAttnProcessor(
-        #     attnstore=controller, place_in_unet=place_in_unet
-        # )
-        attn_procs[name] = CountingCrossAttnProcessor1(
-            attnstore=controller, place_in_unet=place_in_unet
-        )
-
-    model.unet.set_attn_processor(attn_procs)
-    controller.num_att_layers = cross_att_count
-
-def register_hier_output(model):
-    self = model.unet
-    from ldm.modules.diffusionmodules.util import checkpoint, timestep_embedding
-    logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-    def forward(sample, timestep=None, encoder_hidden_states=None, class_labels=None, timestep_cond=None, 
-                attention_mask=None, cross_attention_kwargs=None, added_cond_kwargs=None, down_block_additional_residuals=None,
-                mid_block_additional_residual=None, encoder_attention_mask=None, return_dict=True):
-
-        out_list = []
-
-        
-        default_overall_up_factor = 2**self.num_upsamplers
-
-        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
-        forward_upsample_size = False
-        upsample_size = None
-
-        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
-            logger.info("Forward upsample size to force interpolation output size.")
-            forward_upsample_size = True
-
-        if attention_mask is not None:
-            # assume that mask is expressed as:
-            #   (1 = keep,      0 = discard)
-            # convert mask into a bias that can be added to attention scores:
-            #       (keep = +0,     discard = -10000.0)
-            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
-            attention_mask = attention_mask.unsqueeze(1)
-
-        if encoder_attention_mask is not None:
-            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
-            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
-
-        if self.config.center_input_sample:
-            sample = 2 * sample - 1.0
-
-        timesteps = timestep
-        if not torch.is_tensor(timesteps):
-            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
-            # This would be a good case for the `match` statement (Python 3.10+)
-            is_mps = sample.device.type == "mps"
-            if isinstance(timestep, float):
-                dtype = torch.float32 if is_mps else torch.float64
-            else:
-                dtype = torch.int32 if is_mps else torch.int64
-            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
-        elif len(timesteps.shape) == 0:
-            timesteps = timesteps[None].to(sample.device)
-
-        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-        timesteps = timesteps.expand(sample.shape[0])
-
-        t_emb = self.time_proj(timesteps)
-
-        t_emb = t_emb.to(dtype=sample.dtype)
-
-        emb = self.time_embedding(t_emb, timestep_cond)
-        aug_emb = None
-
-        if self.class_embedding is not None:
-            if class_labels is None:
-                raise ValueError("class_labels should be provided when num_class_embeds > 0")
-
-            if self.config.class_embed_type == "timestep":
-                class_labels = self.time_proj(class_labels)
-
-                # `Timesteps` does not contain any weights and will always return f32 tensors
-                # there might be better ways to encapsulate this.
-                class_labels = class_labels.to(dtype=sample.dtype)
-
-            class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype)
-
-            if self.config.class_embeddings_concat:
-                emb = torch.cat([emb, class_emb], dim=-1)
-            else:
-                emb = emb + class_emb
-
-        if self.config.addition_embed_type == "text":
-            aug_emb = self.add_embedding(encoder_hidden_states)
-        elif self.config.addition_embed_type == "text_image":
-            # Kandinsky 2.1 - style
-            if "image_embeds" not in added_cond_kwargs:
-                raise ValueError(
-                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
-                )
-
-            image_embs = added_cond_kwargs.get("image_embeds")
-            text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states)
-            aug_emb = self.add_embedding(text_embs, image_embs)
-        elif self.config.addition_embed_type == "text_time":
-            # SDXL - style
-            if "text_embeds" not in added_cond_kwargs:
-                raise ValueError(
-                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
-                )
-            text_embeds = added_cond_kwargs.get("text_embeds")
-            if "time_ids" not in added_cond_kwargs:
-                raise ValueError(
-                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
-                )
-            time_ids = added_cond_kwargs.get("time_ids")
-            time_embeds = self.add_time_proj(time_ids.flatten())
-            time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
-
-            add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
-            add_embeds = add_embeds.to(emb.dtype)
-            aug_emb = self.add_embedding(add_embeds)
-        elif self.config.addition_embed_type == "image":
-            # Kandinsky 2.2 - style
-            if "image_embeds" not in added_cond_kwargs:
-                raise ValueError(
-                    f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
-                )
-            image_embs = added_cond_kwargs.get("image_embeds")
-            aug_emb = self.add_embedding(image_embs)
-        elif self.config.addition_embed_type == "image_hint":
-            # Kandinsky 2.2 - style
-            if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs:
-                raise ValueError(
-                    f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`"
-                )
-            image_embs = added_cond_kwargs.get("image_embeds")
-            hint = added_cond_kwargs.get("hint")
-            aug_emb, hint = self.add_embedding(image_embs, hint)
-            sample = torch.cat([sample, hint], dim=1)
-
-        emb = emb + aug_emb if aug_emb is not None else emb
-
-        if self.time_embed_act is not None:
-            emb = self.time_embed_act(emb)
-
-        if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj":
-            encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)
-        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj":
-            # Kadinsky 2.1 - style
-            if "image_embeds" not in added_cond_kwargs:
-                raise ValueError(
-                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
-                )
-
-            image_embeds = added_cond_kwargs.get("image_embeds")
-            encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds)
-        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj":
-            # Kandinsky 2.2 - style
-            if "image_embeds" not in added_cond_kwargs:
-                raise ValueError(
-                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
-                )
-            image_embeds = added_cond_kwargs.get("image_embeds")
-            encoder_hidden_states = self.encoder_hid_proj(image_embeds)
-        # 2. pre-process
-        sample = self.conv_in(sample) # 1, 320, 64, 64
-
-        # 2.5 GLIGEN position net
-        if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None:
-            cross_attention_kwargs = cross_attention_kwargs.copy()
-            gligen_args = cross_attention_kwargs.pop("gligen")
-            cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)}
-
-        # 3. down
-        lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
-
-        is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
-        is_adapter = mid_block_additional_residual is None and down_block_additional_residuals is not None
-
-        down_block_res_samples = (sample,)
-
-        for downsample_block in self.down_blocks:
-            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
-                # For t2i-adapter CrossAttnDownBlock2D
-                additional_residuals = {}
-                if is_adapter and len(down_block_additional_residuals) > 0:
-                    additional_residuals["additional_residuals"] = down_block_additional_residuals.pop(0)
-
-                sample, res_samples = downsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    encoder_hidden_states=encoder_hidden_states,
-                    attention_mask=attention_mask,
-                    cross_attention_kwargs=cross_attention_kwargs,
-                    encoder_attention_mask=encoder_attention_mask,
-                    **additional_residuals,
-                )
-            else:
-                sample, res_samples = downsample_block(hidden_states=sample, temb=emb, scale=lora_scale)
-
-                if is_adapter and len(down_block_additional_residuals) > 0:
-                    sample += down_block_additional_residuals.pop(0)
-
-            down_block_res_samples += res_samples
-
-        if is_controlnet:
-            new_down_block_res_samples = ()
-
-            for down_block_res_sample, down_block_additional_residual in zip(
-                down_block_res_samples, down_block_additional_residuals
-            ):
-                down_block_res_sample = down_block_res_sample + down_block_additional_residual
-                new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
-
-            down_block_res_samples = new_down_block_res_samples
-
-        # 4. mid
-        if self.mid_block is not None:
-            sample = self.mid_block(
-                sample,
-                emb,
-                encoder_hidden_states=encoder_hidden_states,
-                attention_mask=attention_mask,
-                cross_attention_kwargs=cross_attention_kwargs,
-                encoder_attention_mask=encoder_attention_mask,
-            )
-            # To support T2I-Adapter-XL
-            if (
-                is_adapter
-                and len(down_block_additional_residuals) > 0
-                and sample.shape == down_block_additional_residuals[0].shape
-            ):
-                sample += down_block_additional_residuals.pop(0)
-
-        if is_controlnet:
-            sample = sample + mid_block_additional_residual
-
-        # 5. up
-        for i, upsample_block in enumerate(self.up_blocks):
-            is_final_block = i == len(self.up_blocks) - 1
-
-            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
-            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
-
-            # if we have not reached the final block and need to forward the
-            # upsample size, we do it here
-            if not is_final_block and forward_upsample_size:
-                upsample_size = down_block_res_samples[-1].shape[2:]
-
-            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
-                sample = upsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    res_hidden_states_tuple=res_samples,
-                    encoder_hidden_states=encoder_hidden_states,
-                    cross_attention_kwargs=cross_attention_kwargs,
-                    upsample_size=upsample_size,
-                    attention_mask=attention_mask,
-                    encoder_attention_mask=encoder_attention_mask,
-                )
-            else:
-                sample = upsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    res_hidden_states_tuple=res_samples,
-                    upsample_size=upsample_size,
-                    scale=lora_scale,
-                )
-
-            # if i in [1, 4, 7]:
-            out_list.append(sample)
-
-        # 6. post-process
-        if self.conv_norm_out:
-            sample = self.conv_norm_out(sample)
-            sample = self.conv_act(sample)
-        sample = self.conv_out(sample)
-
-        if not return_dict:
-            return (sample,)
-
-        return UNet2DConditionOutput(sample=sample), out_list
-    
-    self.forward = forward
-
-
-class AttentionControl(abc.ABC):
-
-    def step_callback(self, x_t):
-        return x_t
-
-    def between_steps(self):
-        return
-
-    @property
-    def num_uncond_att_layers(self):
-        return 0
-
-    @abc.abstractmethod
-    def forward(self, attn, is_cross: bool, place_in_unet: str):
-        raise NotImplementedError
-
-    def __call__(self, attn, is_cross: bool, place_in_unet: str):
-        if self.cur_att_layer >= self.num_uncond_att_layers:
-            # self.forward(attn, is_cross, place_in_unet)
-            if LOW_RESOURCE:
-                attn = self.forward(attn, is_cross, place_in_unet)
-            else:
-                h = attn.shape[0]
-                attn[h // 2:] = self.forward(attn[h // 2:], is_cross, place_in_unet)
-        self.cur_att_layer += 1
-        if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers:
-            self.cur_att_layer = 0
-            self.cur_step += 1
-            self.between_steps()
-        return attn
-
-    def reset(self):
-        self.cur_step = 0
-        self.cur_att_layer = 0
-
-    def __init__(self):
-        self.cur_step = 0
-        self.num_att_layers = -1
-        self.cur_att_layer = 0
-
-
-class EmptyControl(AttentionControl):
-
-    def forward(self, attn, is_cross: bool, place_in_unet: str):
-        return attn
-
-
-class AttentionStore(AttentionControl):
-
-    @staticmethod
-    def get_empty_store():
-        return {"down_cross": [], "mid_cross": [], "up_cross": [],
-                "down_self": [], "mid_self": [], "up_self": []}
-
-    def forward(self, attn, is_cross: bool, place_in_unet: str):
-        key = f"{place_in_unet}_{'cross' if is_cross else 'self'}"
-        if attn.shape[1] <= self.max_size ** 2:  # avoid memory overhead
-            self.step_store[key].append(attn)
-        return attn
-
-    def between_steps(self):
-        self.attention_store = self.step_store
-        if self.save_global_store:
-            with torch.no_grad():
-                if len(self.global_store) == 0:
-                    self.global_store = self.step_store
-                else:
-                    for key in self.global_store:
-                        for i in range(len(self.global_store[key])):
-                            self.global_store[key][i] += self.step_store[key][i].detach()
-        self.step_store = self.get_empty_store()
-        self.step_store = self.get_empty_store()
-
-    def get_average_attention(self):
-        average_attention = self.attention_store
-        return average_attention
-
-    def get_average_global_attention(self):
-        average_attention = {key: [item / self.cur_step for item in self.global_store[key]] for key in
-                             self.attention_store}
-        return average_attention
-
-    def reset(self):
-        super(AttentionStore, self).reset()
-        self.step_store = self.get_empty_store()
-        self.attention_store = {}
-        self.global_store = {}
-
-    def __init__(self, max_size=32, save_global_store=False):
-        '''
-        Initialize an empty AttentionStore
-        :param step_index: used to visualize only a specific step in the diffusion process
-        '''
-        super(AttentionStore, self).__init__()
-        self.save_global_store = save_global_store
-        self.max_size = max_size
-        self.step_store = self.get_empty_store()
-        self.attention_store = {}
-        self.global_store = {}
-        self.curr_step_index = 0
-
-def aggregate_attention(prompts, attention_store: AttentionStore, res: int, from_where: List[str], is_cross: bool, select: int):
-    out = []
-    attention_maps = attention_store.get_average_attention()
-    num_pixels = res ** 2
-    for location in from_where:
-        for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]:
-            if item.shape[1] == num_pixels:
-                cross_maps = item.reshape(len(prompts), -1, res, res, item.shape[-1])[select]
-                out.append(cross_maps)
-    out = torch.cat(out, dim=0)
-    out = out.sum(0) / out.shape[0]
-    return out
-
-
-def show_cross_attention(tokenizer, prompts, attention_store: AttentionStore, res: int, from_where: List[str], select: int = 0):
-    tokens = tokenizer.encode(prompts[select])
-    decoder = tokenizer.decode
-    attention_maps = aggregate_attention(attention_store, res, from_where, True, select)
-    images = []
-    for i in range(len(tokens)):
-        image = attention_maps[:, :, i]
-        image = 255 * image / image.max()
-        image = image.unsqueeze(-1).expand(*image.shape, 3)
-        image = image.numpy().astype(np.uint8)
-        image = np.array(Image.fromarray(image).resize((256, 256)))
-        image = text_under_image(image, decoder(int(tokens[i])))
-        images.append(image)
-    view_images(np.stack(images, axis=0))
-    
-
-def show_self_attention_comp(attention_store: AttentionStore, res: int, from_where: List[str],
-                        max_com=10, select: int = 0):
-    attention_maps = aggregate_attention(attention_store, res, from_where, False, select).numpy().reshape((res ** 2, res ** 2))
-    u, s, vh = np.linalg.svd(attention_maps - np.mean(attention_maps, axis=1, keepdims=True))
-    images = []
-    for i in range(max_com):
-        image = vh[i].reshape(res, res)
-        image = image - image.min()
-        image = 255 * image / image.max()
-        image = np.repeat(np.expand_dims(image, axis=2), 3, axis=2).astype(np.uint8)
-        image = Image.fromarray(image).resize((256, 256))
-        image = np.array(image)
-        images.append(image)
-    view_images(np.concatenate(images, axis=1))
-
-def text_under_image(image: np.ndarray, text: str, text_color: Tuple[int, int, int] = (0, 0, 0)):
-    h, w, c = image.shape
-    offset = int(h * .2)
-    img = np.ones((h + offset, w, c), dtype=np.uint8) * 255
-    font = cv2.FONT_HERSHEY_SIMPLEX
-    # font = ImageFont.truetype("/usr/share/fonts/truetype/noto/NotoMono-Regular.ttf", font_size)
-    img[:h] = image
-    textsize = cv2.getTextSize(text, font, 1, 2)[0]
-    text_x, text_y = (w - textsize[0]) // 2, h + offset - textsize[1] // 2
-    cv2.putText(img, text, (text_x, text_y ), font, 1, text_color, 2)
-    return img
-
-
-def view_images(images, num_rows=1, offset_ratio=0.02):
-    if type(images) is list:
-        num_empty = len(images) % num_rows
-    elif images.ndim == 4:
-        num_empty = images.shape[0] % num_rows
-    else:
-        images = [images]
-        num_empty = 0
-
-    empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
-    images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
-    num_items = len(images)
-
-    h, w, c = images[0].shape
-    offset = int(h * offset_ratio)
-    num_cols = num_items // num_rows
-    image_ = np.ones((h * num_rows + offset * (num_rows - 1),
-                      w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
-    for i in range(num_rows):
-        for j in range(num_cols):
-            image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[
-                i * num_cols + j]
-
-    pil_img = Image.fromarray(image_)
-    display(pil_img)
-
-def self_cross_attn(self_attn, cross_attn):
-    res = self_attn.shape[0]
-    assert res == cross_attn.shape[0]
-    # cross attn [res, res] -> [res*res]
-    cross_attn_ = cross_attn.reshape([res*res])
-    # self_attn [res, res, res*res]
-    self_cross_attn = cross_attn_ * self_attn
-    self_cross_attn = self_cross_attn.mean(-1).unsqueeze(0).unsqueeze(0)
-    return self_cross_attn

_utils/attn_utils_new.py

@@ -19,7 +19,6 @@ try:
     from diffusers.models.cross_attention import CrossAttention
 except:
     from diffusers.models.attention_processor import Attention as CrossAttention
-from typing import Any, Dict, List, Optional, Tuple, Union
 MAX_NUM_WORDS = 77
 LOW_RESOURCE = False 
 
@@ -512,91 +511,7 @@ def aggregate_attention(prompts, attention_store: AttentionStore, res: int, from
     out = out.sum(0) / out.shape[0]
     return out
 
-def aggregate_attention1(prompts, attention_store: AttentionStore, res: int, from_where: List[str], is_cross: bool, select: int):
-    out = []
-    attention_maps = attention_store.get_average_attention()
-    num_pixels = res ** 2
-    for location in from_where:
-        for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]:
-            if item.shape[1] == num_pixels:
-                cross_maps = item.reshape(len(prompts), -1, res, res, item.shape[-1])[select]
-                out.append(cross_maps)
-    # out = torch.cat(out, dim=0)
-    # out = out.sum(0) / out.shape[0]
-    out = out[1]
-    out = out.sum(0) / out.shape[0]
-    return out
-
-
-def show_cross_attention(tokenizer, prompts, attention_store: AttentionStore, res: int, from_where: List[str], select: int = 0):
-    tokens = tokenizer.encode(prompts[select])
-    decoder = tokenizer.decode
-    attention_maps = aggregate_attention(attention_store, res, from_where, True, select)
-    images = []
-    for i in range(len(tokens)):
-        image = attention_maps[:, :, i]
-        image = 255 * image / image.max()
-        image = image.unsqueeze(-1).expand(*image.shape, 3)
-        image = image.numpy().astype(np.uint8)
-        image = np.array(Image.fromarray(image).resize((256, 256)))
-        image = text_under_image(image, decoder(int(tokens[i])))
-        images.append(image)
-    view_images(np.stack(images, axis=0))
-    
 
-def show_self_attention_comp(attention_store: AttentionStore, res: int, from_where: List[str],
-                        max_com=10, select: int = 0):
-    attention_maps = aggregate_attention(attention_store, res, from_where, False, select).numpy().reshape((res ** 2, res ** 2))
-    u, s, vh = np.linalg.svd(attention_maps - np.mean(attention_maps, axis=1, keepdims=True))
-    images = []
-    for i in range(max_com):
-        image = vh[i].reshape(res, res)
-        image = image - image.min()
-        image = 255 * image / image.max()
-        image = np.repeat(np.expand_dims(image, axis=2), 3, axis=2).astype(np.uint8)
-        image = Image.fromarray(image).resize((256, 256))
-        image = np.array(image)
-        images.append(image)
-    view_images(np.concatenate(images, axis=1))
-
-def text_under_image(image: np.ndarray, text: str, text_color: Tuple[int, int, int] = (0, 0, 0)):
-    h, w, c = image.shape
-    offset = int(h * .2)
-    img = np.ones((h + offset, w, c), dtype=np.uint8) * 255
-    font = cv2.FONT_HERSHEY_SIMPLEX
-    # font = ImageFont.truetype("/usr/share/fonts/truetype/noto/NotoMono-Regular.ttf", font_size)
-    img[:h] = image
-    textsize = cv2.getTextSize(text, font, 1, 2)[0]
-    text_x, text_y = (w - textsize[0]) // 2, h + offset - textsize[1] // 2
-    cv2.putText(img, text, (text_x, text_y ), font, 1, text_color, 2)
-    return img
-
-
-def view_images(images, num_rows=1, offset_ratio=0.02):
-    if type(images) is list:
-        num_empty = len(images) % num_rows
-    elif images.ndim == 4:
-        num_empty = images.shape[0] % num_rows
-    else:
-        images = [images]
-        num_empty = 0
-
-    empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
-    images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
-    num_items = len(images)
-
-    h, w, c = images[0].shape
-    offset = int(h * offset_ratio)
-    num_cols = num_items // num_rows
-    image_ = np.ones((h * num_rows + offset * (num_rows - 1),
-                      w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
-    for i in range(num_rows):
-        for j in range(num_cols):
-            image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[
-                i * num_cols + j]
-
-    pil_img = Image.fromarray(image_)
-    display(pil_img)
 
 def self_cross_attn(self_attn, cross_attn):
     cross_attn = cross_attn.squeeze()

_utils/load_models.py

@@ -6,11 +6,7 @@ import torch.nn as nn
 def load_stable_diffusion_model(config: RunConfig):
     device = torch.device('cpu')
 
-    if config.sd_2_1:
-        stable_diffusion_version = "stabilityai/stable-diffusion-2-1-base"
-    else:
-        stable_diffusion_version = "CompVis/stable-diffusion-v1-4"
-    # stable = StableCountingPipeline.from_pretrained(stable_diffusion_version).to(device)
+    stable_diffusion_version = "CompVis/stable-diffusion-v1-4"
     stable = StableDiffusionPipeline.from_pretrained(stable_diffusion_version).to(device)
     return stable
 

_utils/seg_eval.py

@@ -1,61 +0,0 @@
-import torch
-
-
-def iou_torch(inst1, inst2):
-    inter = torch.logical_and(inst1, inst2).sum().float()
-    union = torch.logical_or(inst1, inst2).sum().float()
-    if union == 0:
-        return torch.tensor(float('nan'))
-    return inter / union
-
-def get_instances_torch(mask):
-    # 返回所有非背景的 instance mask（布尔型）
-    ids = torch.unique(mask)
-    return [(mask == i) for i in ids if i != 0]
-
-def compute_instance_miou(pred_mask, gt_mask):
-    # pred_mask 和 gt_mask 都是 torch.Tensor, shape [H, W], 整数类型
-    pred_instances = get_instances_torch(pred_mask)
-    gt_instances = get_instances_torch(gt_mask)
-
-    ious = []
-    for gt in gt_instances:
-        best_iou = torch.tensor(0.0).to(pred_mask.device)
-        for pred in pred_instances:
-            i = iou_torch(pred, gt)
-            if i > best_iou:
-                best_iou = i
-        ious.append(best_iou)
-    
-    # 处理空情况
-    if len(ious) == 0:
-        return torch.tensor(float('nan'))
-    return torch.nanmean(torch.stack(ious))
-
-from torch import Tensor
-
-
-def dice_coeff(input: Tensor, target: Tensor, reduce_batch_first: bool = False, epsilon: float = 1e-6):
-    # Average of Dice coefficient for all batches, or for a single mask
-    assert input.size() == target.size()
-    assert input.dim() == 3 or not reduce_batch_first
-
-    sum_dim = (-1, -2) if input.dim() == 2 or not reduce_batch_first else (-1, -2, -3)
-
-    inter = 2 * (input * target).sum(dim=sum_dim)
-    sets_sum = input.sum(dim=sum_dim) + target.sum(dim=sum_dim)
-    sets_sum = torch.where(sets_sum == 0, inter, sets_sum)
-
-    dice = (inter + epsilon) / (sets_sum + epsilon)
-    return dice.mean()
-
-
-def multiclass_dice_coeff(input: Tensor, target: Tensor, reduce_batch_first: bool = False, epsilon: float = 1e-6):
-    # Average of Dice coefficient for all classes
-    return dice_coeff(input.flatten(0, 1), target.flatten(0, 1), reduce_batch_first, epsilon)
-
-
-def dice_loss(input: Tensor, target: Tensor, multiclass: bool = False):
-    # Dice loss (objective to minimize) between 0 and 1
-    fn = multiclass_dice_coeff if multiclass else dice_coeff
-    return 1 - fn(input, target, reduce_batch_first=True)

config.py

@@ -7,8 +7,6 @@ from typing import Dict, List
 class RunConfig:
     # Guiding text prompt
     prompt: str = "<task-prompt>"
-    # Whether to use Stable Diffusion v2.1
-    sd_2_1: bool = False
     # Which token indices to alter with attend-and-excite
     token_indices: List[int] = field(default_factory=lambda: [2,5])
     # Which random seeds to use when generating

counting.py

@@ -12,19 +12,16 @@ from PIL import Image
 import numpy as np
 from config import RunConfig
 from _utils import attn_utils_new as attn_utils
-from _utils.attn_utils import AttentionStore
+from _utils.attn_utils_new import AttentionStore
 from _utils.misc_helper import *
 import torch.nn.functional as F
 import matplotlib.pyplot as plt
 import cv2
 import warnings
-from pytorch_lightning.callbacks import ModelCheckpoint
 warnings.filterwarnings("ignore", category=UserWarning)
 import pytorch_lightning as pl
 from _utils.load_models import load_stable_diffusion_model
 from models.model import Counting_with_SD_features_loca as Counting
-from pytorch_lightning.loggers import WandbLogger
-from models.enc_model.loca_args import get_argparser as loca_get_argparser
 from models.enc_model.loca import build_model as build_loca_model
 import time
 import torchvision.transforms as T
@@ -44,12 +41,7 @@ class CountingModule(pl.LightningModule):
     def initialize_model(self):
         
         # load loca model
-        loca_args = loca_get_argparser().parse_args()
-        self.loca_model = build_loca_model(loca_args)
-        # weights = torch.load("ckpt/loca_few_shot.pt")["model"]
-        # weights = {k.replace("module","") : v for k, v in weights.items()}
-        # self.loca_model.load_state_dict(weights, strict=False)
-        # del weights
+        self.loca_model = build_loca_model()
 
         self.counting_adapter = Counting(scale_factor=SCALE)
         # if os.path.isfile(self.args.adapter_weight):

models/enc_model/loca.py

@@ -78,12 +78,6 @@ class LOCA(nn.Module):
             nn.LayerNorm((64, 64))
         )
 
-        # self.fuse1 = nn.Sequential(
-        #     nn.Conv2d(322, 256, kernel_size=1, stride=1),
-        #     nn.LeakyReLU(),
-        #     nn.LayerNorm((64, 64))
-        # )
-
     def forward_before_reg(self, x, bboxes):
         num_objects = bboxes.size(1) if not self.zero_shot else self.num_objects
         # backbone
@@ -105,7 +99,6 @@ class LOCA(nn.Module):
 
         all_prototypes = self.ope(f_e, pos_emb, bboxes) # [3, 27, 1, 256]
 
-        outputs = list()
         response_maps_list = []
         for i in range(all_prototypes.size(0)):
             prototypes = all_prototypes[i, ...].permute(1, 0, 2).reshape(
@@ -122,18 +115,10 @@ class LOCA(nn.Module):
                 bs, num_objects, self.emb_dim, h, w
             ).max(dim=1)[0]
 
-            # # send through regression heads
-            # if i == all_prototypes.size(0) - 1:
-            #     predicted_dmaps = self.regression_head(response_maps)
-            # else:
-            #     predicted_dmaps = self.aux_heads[i](response_maps)
-            # outputs.append(predicted_dmaps)
             response_maps_list.append(response_maps)
 
         out = {
-            # "pred": outputs[-1],
             "feature_bf_regression": response_maps_list[-1],
-            # "aux_pred": outputs[:-1],
             "aux_feature_bf_regression": response_maps_list[:-1]
         }
 
@@ -162,71 +147,61 @@ class LOCA(nn.Module):
         
         return {"pred": outputs[-1], "aux_pred": outputs[:-1]}
     
-    def forward_reg1(self, response_maps, self_attn):
-        # attn_stack = self.attn_norm(attn_stack)
-        # attn_stack_mean = torch.mean(attn_stack, dim=1, keepdim=True)
-        # unet_feature = torch.cat([unet_feature, attn_stack], dim=1) # [1, 324, 64, 64]
-        # unet_feature = unet_feature * attn_stack_mean
-        # if unet_feature.shape[1] == 322:
-        #     unet_feature = self.fuse1(unet_feature)
-        # else:
-        #     unet_feature = self.fuse(unet_feature)
-
-        
-
-        response_maps = response_maps["aux_feature_bf_regression"] + [response_maps["feature_bf_regression"]]
-
-        outputs = []
-        for i in range(len(response_maps)):
-            response_map = response_maps[i] + self_attn
-            if i == len(response_maps) - 1:
-                predicted_dmaps = self.regression_head(response_map)
-            else:
-                predicted_dmaps = self.aux_heads[i](response_map)
-            outputs.append(predicted_dmaps)
+    # def forward_reg1(self, response_maps, self_attn):        
+
+    #     response_maps = response_maps["aux_feature_bf_regression"] + [response_maps["feature_bf_regression"]]
+
+    #     outputs = []
+    #     for i in range(len(response_maps)):
+    #         response_map = response_maps[i] + self_attn
+    #         if i == len(response_maps) - 1:
+    #             predicted_dmaps = self.regression_head(response_map)
+    #         else:
+    #             predicted_dmaps = self.aux_heads[i](response_map)
+    #         outputs.append(predicted_dmaps)
         
-        return {"pred": outputs[-1], "aux_pred": outputs[:-1]}
+    #     return {"pred": outputs[-1], "aux_pred": outputs[:-1]}
     
-    def forward_reg_without_unet(self, response_maps, attn_stack):
-        # attn_stack = self.attn_norm(attn_stack)
-        attn_stack_mean = torch.mean(attn_stack, dim=1, keepdim=True)
-
-        response_maps = response_maps["aux_feature_bf_regression"] + [response_maps["feature_bf_regression"]]
-
-        outputs = []
-        for i in range(len(response_maps)):
-            response_map = response_maps[i] * attn_stack_mean * 0.5 + response_maps[i]
-            if i == len(response_maps) - 1:
-                predicted_dmaps = self.regression_head(response_map)
-            else:
-                predicted_dmaps = self.aux_heads[i](response_map)
-            outputs.append(predicted_dmaps)
+    # def forward_reg_without_unet(self, response_maps, attn_stack):
+    #     # attn_stack = self.attn_norm(attn_stack)
+    #     attn_stack_mean = torch.mean(attn_stack, dim=1, keepdim=True)
+
+    #     response_maps = response_maps["aux_feature_bf_regression"] + [response_maps["feature_bf_regression"]]
+
+    #     outputs = []
+    #     for i in range(len(response_maps)):
+    #         response_map = response_maps[i] * attn_stack_mean * 0.5 + response_maps[i]
+    #         if i == len(response_maps) - 1:
+    #             predicted_dmaps = self.regression_head(response_map)
+    #         else:
+    #             predicted_dmaps = self.aux_heads[i](response_map)
+    #         outputs.append(predicted_dmaps)
         
-        return {"pred": outputs[-1], "aux_pred": outputs[:-1]}
-
-
-def build_model(args):
+    #     return {"pred": outputs[-1], "aux_pred": outputs[:-1]}
 
-    assert args.backbone in ['resnet18', 'resnet50', 'resnet101']
-    assert args.reduction in [4, 8, 16]
 
+def build_model():
+    """
+    Build LOCA with a fixed configuration based on defaults in `loca_args.py`.
+    The `args` parameter is accepted for backward compatibility but ignored.
+    """
     return LOCA(
-        image_size=args.image_size,
-        num_encoder_layers=args.num_enc_layers,
-        num_ope_iterative_steps=args.num_ope_iterative_steps,
-        num_objects=args.num_objects,
-        zero_shot=args.zero_shot,
-        emb_dim=args.emb_dim,
-        num_heads=args.num_heads,
-        kernel_dim=args.kernel_dim,
-        backbone_name=args.backbone,
-        swav_backbone=args.swav_backbone,
-        train_backbone=args.backbone_lr > 0,
-        reduction=args.reduction,
-        dropout=args.dropout,
+        image_size=512,
+        num_encoder_layers=3,
+        num_ope_iterative_steps=3,
+        num_objects=3,
+        zero_shot=False,
+        emb_dim=256,
+        num_heads=8,
+        kernel_dim=3,
+        backbone_name="resnet50",
+        swav_backbone=True,
+        train_backbone=False,  # backbone_lr default is 0 in loca_args.py
+        reduction=8,
+        dropout=0.1,
         layer_norm_eps=1e-5,
         mlp_factor=8,
-        norm_first=args.pre_norm,
+        norm_first=True,
         activation=nn.GELU,
         norm=True,
     )

models/enc_model/loca_args.py

@@ -1,44 +0,0 @@
-import argparse
-
-
-def get_argparser():
-
-    parser = argparse.ArgumentParser("LOCA parser", add_help=False)
-
-    parser.add_argument('--model_name', default='loca_few_shot', type=str)
-    parser.add_argument(
-        '--data_path',
-        default='./data/FSC147_384_V2',
-        type=str
-    )
-    parser.add_argument(
-        '--model_path',
-        default='ckpt',
-        type=str
-    )
-    parser.add_argument('--backbone', default='resnet50', type=str)
-    parser.add_argument('--swav_backbone', action='store_true', default=True)
-    parser.add_argument('--reduction', default=8, type=int)
-    parser.add_argument('--image_size', default=512, type=int)
-    parser.add_argument('--num_enc_layers', default=3, type=int)
-    parser.add_argument('--num_ope_iterative_steps', default=3, type=int)
-    parser.add_argument('--emb_dim', default=256, type=int)
-    parser.add_argument('--num_heads', default=8, type=int)
-    parser.add_argument('--kernel_dim', default=3, type=int)
-    parser.add_argument('--num_objects', default=3, type=int)
-    parser.add_argument('--epochs', default=200, type=int)
-    parser.add_argument('--resume_training', action='store_true')
-    parser.add_argument('--lr', default=1e-4, type=float)
-    parser.add_argument('--backbone_lr', default=0, type=float)
-    parser.add_argument('--lr_drop', default=200, type=int)
-    parser.add_argument('--weight_decay', default=1e-4, type=float)
-    parser.add_argument('--batch_size', default=1, type=int)
-    parser.add_argument('--dropout', default=0.1, type=float)
-    parser.add_argument('--num_workers', default=8, type=int)
-    parser.add_argument('--max_grad_norm', default=0.1, type=float)
-    parser.add_argument('--aux_weight', default=0.3, type=float)
-    parser.add_argument('--tiling_p', default=0.5, type=float)
-    parser.add_argument('--zero_shot', action='store_true')
-    parser.add_argument('--pre_norm', action='store_true', default=True)
-
-    return parser

models/enc_model/regression_head.py

@@ -55,38 +55,3 @@ class DensityMapRegressor(nn.Module):
                     nn.init.constant_(module.bias, 0)
 
 
-class DensityMapRegressor_(nn.Module):
-
-    def __init__(self, in_channels, reduction):
-
-        super(DensityMapRegressor, self).__init__()
-
-        if reduction == 8:
-            self.regressor = nn.Sequential(
-                UpsamplingLayer(in_channels, 128),
-                UpsamplingLayer(128, 64),
-                UpsamplingLayer(64, 32),
-                nn.Conv2d(32, 1, kernel_size=1),
-                nn.LeakyReLU()
-            )
-        elif reduction == 16:
-            self.regressor = nn.Sequential(
-                UpsamplingLayer(in_channels, 128),
-                UpsamplingLayer(128, 64),
-                UpsamplingLayer(64, 32),
-                UpsamplingLayer(32, 16),
-                nn.Conv2d(16, 1, kernel_size=1),
-                nn.LeakyReLU()
-            )
-
-        self.reset_parameters()
-
-    def forward(self, x):
-        return self.regressor(x)
-
-    def reset_parameters(self):
-        for module in self.modules():
-            if isinstance(module, nn.Conv2d):
-                nn.init.normal_(module.weight, std=0.01)
-                if module.bias is not None:
-                    nn.init.constant_(module.bias, 0)

models/tra_post_model/trackastra/model/model.py

@@ -139,138 +139,6 @@ class DecoderLayer(nn.Module):
         return x
 
 
-# class BidirectionalRelativePositionalAttention(RelativePositionalAttention):
-#     def forward(
-#         self,
-#         query1: torch.Tensor,
-#         query2: torch.Tensor,
-#         coords: torch.Tensor,
-#         padding_mask: torch.Tensor = None,
-#     ):
-#         B, N, D = query1.size()
-#         q1 = self.q_pro(query1)  # (B, N, D)
-#         q2 = self.q_pro(query2)  # (B, N, D)
-#         v1 = self.v_pro(query1)  # (B, N, D)
-#         v2 = self.v_pro(query2)  # (B, N, D)
-
-#         # (B, nh, N, hs)
-#         q1 = q1.view(B, N, self.n_head, D // self.n_head).transpose(1, 2)
-#         v1 = v1.view(B, N, self.n_head, D // self.n_head).transpose(1, 2)
-#         q2 = q2.view(B, N, self.n_head, D // self.n_head).transpose(1, 2)
-#         v2 = v2.view(B, N, self.n_head, D // self.n_head).transpose(1, 2)
-
-#         attn_mask = torch.zeros(
-#             (B, self.n_head, N, N), device=query1.device, dtype=q1.dtype
-#         )
-
-#         # add negative value but not too large to keep mixed precision loss from becoming nan
-#         attn_ignore_val = -1e3
-
-#         # spatial cutoff
-#         yx = coords[..., 1:]
-#         spatial_dist = torch.cdist(yx, yx)
-#         spatial_mask = (spatial_dist > self.cutoff_spatial).unsqueeze(1)
-#         attn_mask.masked_fill_(spatial_mask, attn_ignore_val)
-
-#         # dont add positional bias to self-attention if coords is None
-#         if coords is not None:
-#             if self._mode == "bias":
-#                 attn_mask = attn_mask + self.pos_bias(coords)
-#             elif self._mode == "rope":
-#                 q1, q2 = self.rot_pos_enc(q1, q2, coords)
-#             else:
-#                 pass
-
-#             dist = torch.cdist(coords, coords, p=2)
-#             attn_mask += torch.exp(-0.1 * dist.unsqueeze(1))
-
-#         # if given key_padding_mask = (B,N) then ignore those tokens (e.g. padding tokens)
-#         if padding_mask is not None:
-#             ignore_mask = torch.logical_or(
-#                 padding_mask.unsqueeze(1), padding_mask.unsqueeze(2)
-#             ).unsqueeze(1)
-#             attn_mask.masked_fill_(ignore_mask, attn_ignore_val)
-
-#         self.attn_mask = attn_mask.clone()
-
-#         y1 = nn.functional.scaled_dot_product_attention(
-#             q1,
-#             q2,
-#             v1,
-#             attn_mask=attn_mask,
-#             dropout_p=self.dropout if self.training else 0,
-#         )
-#         y2 = nn.functional.scaled_dot_product_attention(
-#             q2,
-#             q1,
-#             v2,
-#             attn_mask=attn_mask,
-#             dropout_p=self.dropout if self.training else 0,
-#         )
-
-#         y1 = y1.transpose(1, 2).contiguous().view(B, N, D)
-#         y1 = self.proj(y1)
-#         y2 = y2.transpose(1, 2).contiguous().view(B, N, D)
-#         y2 = self.proj(y2)
-#         return y1, y2
-
-
-# class BidirectionalCrossAttention(nn.Module):
-#     def __init__(
-#         self,
-#         coord_dim: int = 2,
-#         d_model=256,
-#         num_heads=4,
-#         dropout=0.1,
-#         window: int = 16,
-#         cutoff_spatial: int = 256,
-#         positional_bias: Literal["bias", "rope", "none"] = "bias",
-#         positional_bias_n_spatial: int = 32,
-#     ):
-#         super().__init__()
-#         self.positional_bias = positional_bias
-#         self.attn = BidirectionalRelativePositionalAttention(
-#             coord_dim,
-#             d_model,
-#             num_heads,
-#             cutoff_spatial=cutoff_spatial,
-#             n_spatial=positional_bias_n_spatial,
-#             cutoff_temporal=window,
-#             n_temporal=window,
-#             dropout=dropout,
-#             mode=positional_bias,
-#         )
-
-#         self.mlp = FeedForward(d_model)
-#         self.norm1 = nn.LayerNorm(d_model)
-#         self.norm2 = nn.LayerNorm(d_model)
-
-#     def forward(
-#         self,
-#         x: torch.Tensor,
-#         y: torch.Tensor,
-#         coords: torch.Tensor,
-#         padding_mask: torch.Tensor = None,
-#     ):
-#         x = self.norm1(x)
-#         y = self.norm1(y)
-
-#         # cross attention
-#         # setting coords to None disables positional bias
-#         x2, y2 = self.attn(
-#             x,
-#             y,
-#             coords=coords if self.positional_bias else None,
-#             padding_mask=padding_mask,
-#         )
-#         # print(torch.norm(x2).item()/torch.norm(x).item())
-#         x = x + x2
-#         x = x + self.mlp(self.norm2(x))
-#         y = y + y2
-#         y = y + self.mlp(self.norm2(y))
-
-#         return x, y
-
 
 class TrackingTransformer(torch.nn.Module):
     def __init__(

segmentation.py

@@ -8,7 +8,7 @@ from PIL import Image
 import numpy as np
 from config import RunConfig
 from _utils import attn_utils_new as attn_utils
-from _utils.attn_utils import AttentionStore
+from _utils.attn_utils_new import AttentionStore
 from _utils.misc_helper import *
 import torch.nn.functional as F
 import logging
@@ -20,10 +20,8 @@ warnings.filterwarnings("ignore", category=UserWarning)
 import pytorch_lightning as pl
 from _utils.load_models import load_stable_diffusion_model
 from models.model import Counting_with_SD_features_dino_vit_c3 as Counting
-from models.enc_model.loca_args import get_argparser as loca_get_argparser
 from models.enc_model.loca import build_model as build_loca_model
 import time
-from _utils.seg_eval import *
 from models.seg_post_model import metrics
 from datetime import datetime
 import json
@@ -49,8 +47,7 @@ class SegmentationModule(pl.LightningModule):
     def initialize_model(self):
         
         # load loca model
-        loca_args = loca_get_argparser().parse_args()
-        self.loca_model = build_loca_model(loca_args)
+        self.loca_model = build_loca_model()
         self.loca_model.eval()
 
         self.counting_adapter = Counting(scale_factor=SCALE)

tracking_one.py

@@ -13,13 +13,9 @@ import tifffile
 import skimage.io as io
 from config import RunConfig
 from _utils import attn_utils_new as attn_utils
-from _utils.attn_utils import AttentionStore
+from _utils.attn_utils_new import AttentionStore
 from _utils.misc_helper import *
-from torch.autograd import Variable
-import itertools
-from accelerate import Accelerator
 import torch.nn.functional as F
-from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline
 from tqdm import tqdm
 import torch.nn as nn
 import matplotlib.pyplot as plt
@@ -29,19 +25,14 @@ warnings.filterwarnings("ignore", category=UserWarning)
 import pytorch_lightning as pl
 from _utils.load_models import load_stable_diffusion_model
 from models.model import Counting_with_SD_features_track as Counting
-from models.enc_model.loca_args import get_argparser as loca_get_argparser
 from models.enc_model.loca import build_model as build_loca_model
 import time
-from _utils.seg_eval import *
-from models.tra_post_model.trackastra.model import Trackastra
 from models.tra_post_model.trackastra.model import TrackingTransformer
 from models.tra_post_model.trackastra.utils import (
-    blockwise_causal_norm,
-    blockwise_sum,
     normalize,
 )
-from models.tra_post_model.trackastra.data import build_windows_sd, get_features, load_tiff_timeseries
-from models.tra_post_model.trackastra.tracking import TrackGraph, build_graph, track_greedy, graph_to_ctc
+from models.tra_post_model.trackastra.data import build_windows_sd, get_features
+from models.tra_post_model.trackastra.tracking import TrackGraph, build_graph, track_greedy
 from _utils.track_args import parse_train_args as get_track_args
 import torchvision.transforms as T
 from pathlib import Path
@@ -90,8 +81,7 @@ class TrackingModule(pl.LightningModule):
     def initialize_model(self):
         
         # load loca model
-        loca_args = loca_get_argparser().parse_args()
-        self.loca_model = build_loca_model(loca_args)
+        self.loca_model = build_loca_model()
         # weights = torch.load("ckpt/loca_few_shot.pt")["model"]
         # weights = {k.replace("module","") : v for k, v in weights.items()}
         # self.loca_model.load_state_dict(weights, strict=False)
@@ -985,7 +975,6 @@ class TrackingModule(pl.LightningModule):
 
         self.eval()
         imgs, imgs_raw, images_stable, tra_imgs, imgs_01, height, width = load_track_images(file_dir)
-        # tra_imgs = torch.from_numpy(imgs_).float().to(self.device)
         imgs_stable = torch.from_numpy(images_stable).float().to(self.device)
         imgs_enc = torch.from_numpy(imgs).float().to(self.device)
 
@@ -1032,37 +1021,31 @@ class TrackingModule(pl.LightningModule):
         )
         track_graph = self._track_from_predictions(predictions, mode=mode, **kwargs)
         
-        # ctc_tracks, masks_tracked = graph_to_ctc(
-        #     track_graph,
-        #     masks,
-        #     outdir=f"tracked/{dataname}",
-        # )
-
         return track_graph, masks
 
         
 
-def inference(data_path, box=None):
-    if box is not None:
-        use_box = True
-    else:
-        use_box = False
+# def inference(data_path, box=None):
+#     if box is not None:
+#         use_box = True
+#     else:
+#         use_box = False
 
-    model = TrackingModule(use_box=use_box)
-    load_msg = model.load_state_dict(torch.load("pretrained/microscopy_matching_tra.pth"), strict=True)
+#     model = TrackingModule(use_box=use_box)
+#     load_msg = model.load_state_dict(torch.load("pretrained/microscopy_matching_tra.pth"), strict=True)
 
-    model.move_to_device(torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu'))
+#     model.move_to_device(torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu'))
 
 
-    track_graph, masks = model.track(file_dir=data_path, dataname="inference_sequence")
+#     track_graph, masks = model.track(file_dir=data_path, dataname="inference_sequence")
 
-    if not os.path.exists(f"tracked_ours_seg_pred3/"):
-        os.makedirs(f"tracked_ours_seg_pred3/")
-    ctc_tracks, masks_tracked = graph_to_ctc(
-        track_graph,
-        masks,
-        outdir=f"tracked_ours_seg_pred3/",
-    )
+#     if not os.path.exists(f"tracked_ours_seg_pred3/"):
+#         os.makedirs(f"tracked_ours_seg_pred3/")
+#     ctc_tracks, masks_tracked = graph_to_ctc(
+#         track_graph,
+#         masks,
+#         outdir=f"tracked_ours_seg_pred3/",
+#     )
 
-if __name__ == "__main__":
-    inference(data_path="example_imgs/2D+Time/Fluo-N2DL-HeLa/train/Fluo-N2DL-HeLa/02")
+# if __name__ == "__main__":
+#     inference(data_path="example_imgs/2D+Time/Fluo-N2DL-HeLa/train/Fluo-N2DL-HeLa/02")