class_docstr: WaveAttention
class_docstr_tok: ["WaveAttention"]
code: def __init__(self, dim, num_heads, sr_ratio):
      super().__init__()
      self.num_heads = num_heads
      head_dim = dim // num_heads
      self.scale = head_dim**-0.5
      self.sr_ratio = sr_ratio
  
      self.dwt = DWT2D(wave="haar")
      self.i_dwt = IDWT2D(wave="haar")
      self.reduce = nn.Sequential(
          nn.Conv2d(dim, dim // 4, kernel_size=1, padding=0, stride=1),
          nn.BatchNorm2d(dim // 4),
          nn.ReLU(inplace=True),
      )
      self.filter = nn.Sequential(
          nn.Conv2d(
              dim, dim, kernel_size=3, padding=1, stride=1, groups=1
          ),
          nn.BatchNorm2d(dim),
          nn.ReLU(inplace=True),
      )
      self.kv_embed = (
          nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)
          if sr_ratio > 1
          else nn.Identity()
      )
      self.q = nn.Linear(dim, dim)
      self.kv = nn.Sequential(
          nn.LayerNorm(dim), nn.Linear(dim, dim * 2)
      )
      self.proj = nn.Linear(dim + dim // 4, dim)
      self.apply(init_vit_weights)
  
imports: import torch
  from torch import nn
  from timm.models.layers import DropPath
  import numpy
  from towhee.models.wave_vit.wave_vit_utils import DWT2D, IDWT2D
  from towhee.models.utils.init_vit_weights import init_vit_weights
  from towhee.models.layers.mlp import Mlp
label_desc: Use image node_id 1 to create a new WaveAttention object from inherited base classes: nn with example: obj = WaveAttention(dim, num_heads, sr_ratio)