code: def count_flops_attn(model, _x, y):
      """
      A counter for the `thop` package to count the operations in an
      attention operation.
      Meant to be used like:
          macs, params = thop.profile(
              model,
              inputs=(inputs, timestamps),
              custom_ops={QKVAttention: QKVAttention.count_flops},
          )
      """
      b, c, *spatial = y[0].shape
      num_spatial = int(np.prod(spatial))
      # We perform two matmuls with the same number of ops.
      # The first computes the weight matrix, the second computes
      # the combination of the value vectors.
      matmul_ops = 2 * b * (num_spatial**2) * c
      model.total_ops += th.DoubleTensor([matmul_ops])
  
imports: from abc import abstractmethod
  from functools import partial
  import math
  from typing import Iterable
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from ldm.modules.diffusionmodules.util import checkpoint, conv_nd, linear, avg_pool_nd, zero_module, normalization, timestep_embedding
  from ldm.modules.attention import SpatialTransformer
label_desc: Use image node_id 3 for calling a global function with example usage: count_flops_attn(model, _x, y) without return types