text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  from utils.masking import TriangularCausalMask, ProbMask
  from models.encoder import Encoder, EncoderLayer, ConvLayer, EncoderStack
  from models.decoder import Decoder, DecoderLayer
  from models.attn import FullAttention, ProbAttention, AttentionLayer
  from models.embed import DataEmbedding
  
  from models.Conformer_encoder.conformer_encoder import ConformerEncoder
  
  
  class Informer(nn.Module):
      def __init__(self, enc_in, dec_in, c_out, seq_len, label_len, out_len, 
                  factor=5, d_model=512, n_heads=8, e_layers=3, d_layers=2, d_ff=512, 
                  dropout=0.0, attn='prob', embed='fixed', freq='h', activation='gelu', 
                  output_attention = False, distil=True, 
                  device: torch.device = 'cuda'):
          super(Informer, self).__init__()
          self.pred_len = out_len
          self.attn = attn
          self.output_attention = output_attention
  
          # embedding
          self.enc_embedding = DataEmbedding(enc_in, d_model, embed, freq, dropout)
          self.dec_embedding = DataEmbedding(dec_in, d_model, embed, freq, dropout)
  
          # Attention
          Attn = ProbAttention if attn=='prob' else FullAttention
          # Encoder
          # self.encoder = Encoder(
          #     [
          #         EncoderLayer(
          #             AttentionLayer(Attn(False, factor, attention_dropout=dropout, output_attention=output_attention), 
          #                         d_model, n_heads),
          #             d_model,
          #             d_ff,
          #             dropout=dropout,
          #             activation=activation
          #         )
          #          for l in range(e_layers)
          #     ],
          #     [
          #         ConvLayer(
          #             d_model
          #         ) for l in range(e_layers-1)
          #     ] if distil else None,
          #     norm_layer=torch.nn.LayerNorm(d_model)
          # )
          self.encoder = ConformerEncoder(
              encoder_dim=d_model,
              num_layers=e_layers,
              num_attention_heads=n_heads,
              feed_forward_expansion_factor=4,
              conv_expansion_factor=2,
              feed_forward_dropout_p=dropout,
              attention_dropout_p=dropout,
              conv_dropout_p=dropout,
              conv_kernel_size=31,
              half_step_residual=True,
              device=device,
          )
          # # Decoder
          self.decoder = Decoder(
              [
                  DecoderLayer(
                      AttentionLayer(Attn(True, factor, attention_dropout=dropout, output_attention=False), 
                                  d_model, n_heads),
                      AttentionLayer(FullAttention(False, factor, attention_dropout=dropout, output_attention=False), 
                                  d_model, n_heads),
                      d_model,
                      d_ff,
                      dropout=dropout,
                      activation=activation,
                  )
                  for l in range(d_layers)
              ],
              norm_layer=torch.nn.LayerNorm(d_model)
          )
          # self.end_conv1 = nn.Conv1d(in_channels=label_len+out_len, out_channels=out_len, kernel_size=1, bias=True)
          # self.end_conv2 = nn.Conv1d(in_channels=d_model, out_channels=c_out, kernel_size=1, bias=True)
          self.projection = nn.Linear(d_model, c_out, bias=True)
          
      def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, 
                  enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
  
          enc_out = self.enc_embedding(x_enc, x_mark_enc)#enc_out[32, 96, 512]
          # enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
          enc_out= self.encoder(enc_out)
  
          dec_out = self.dec_embedding(x_dec, x_mark_dec)
          dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
          dec_out = self.projection(dec_out) 
          
          # dec_out = self.end_conv1(dec_out)
          # dec_out = self.end_conv2(dec_out.transpose(2,1)).transpose(1,2)
          if self.output_attention:
              # return dec_out[:,-self.pred_len:,:], attns
               return dec_out[:,-self.pred_len:,:]
          else:
              return dec_out[:,-self.pred_len:,:] # [B, L, D]
  
  
  class InformerStack(nn.Module):
      def __init__(self, enc_in, dec_in, c_out, seq_len, label_len, out_len, 
                  factor=5, d_model=512, n_heads=8, e_layers=[3,2,1], d_layers=2, d_ff=512, 
                  dropout=0.0, attn='prob', embed='fixed', freq='h', activation='gelu',
                  output_attention = False, distil=True,
                  device=torch.device('cuda:0')):
          super(InformerStack, self).__init__()
          self.pred_len = out_len
          self.attn = attn
          self.output_attention = output_attention
  
          # Encoding
          self.enc_embedding = DataEmbedding(enc_in, d_model, embed, freq, dropout)
          self.dec_embedding = DataEmbedding(dec_in, d_model, embed, freq, dropout)
          # Attention
          Attn = ProbAttention if attn=='prob' else FullAttention
          # Encoder
  
          inp_lens = list(range(len(e_layers))) # [0,1,2,...] you can customize here
          encoders = [
              Encoder(
                  [
                      EncoderLayer(
                          AttentionLayer(Attn(False, factor, attention_dropout=dropout, output_attention=output_attention), 
                                      d_model, n_heads),
                          d_model,
                          d_ff,
                          dropout=dropout,
                          activation=activation
                      ) for l in range(el)
                  ],
                  [
                      ConvLayer(
                          d_model
                      ) for l in range(el-1)
                  ] if distil else None,
                  norm_layer=torch.nn.LayerNorm(d_model)
              ) for el in e_layers]
          self.encoder = EncoderStack(encoders, inp_lens)
          # Decoder
          self.decoder = Decoder(
              [
                  DecoderLayer(
                      AttentionLayer(Attn(True, factor, attention_dropout=dropout, output_attention=False), 
                                  d_model, n_heads),
                      AttentionLayer(FullAttention(False, factor, attention_dropout=dropout, output_attention=False), 
                                  d_model, n_heads),
                      d_model,
                      d_ff,
                      dropout=dropout,
                      activation=activation,
                  )
                  for l in range(d_layers)
              ],
              norm_layer=torch.nn.LayerNorm(d_model)
          )
          # self.end_conv1 = nn.Conv1d(in_channels=label_len+out_len, out_channels=out_len, kernel_size=1, bias=True)
          # self.end_conv2 = nn.Conv1d(in_channels=d_model, out_channels=c_out, kernel_size=1, bias=True)
          self.projection = nn.Linear(d_model, c_out, bias=True)
          
      def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, 
                  enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
          enc_out = self.enc_embedding(x_enc, x_mark_enc)
          enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
  
          dec_out = self.dec_embedding(x_dec, x_mark_dec)
          dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
          dec_out = self.projection(dec_out)
          
          # dec_out = self.end_conv1(dec_out)
          # dec_out = self.end_conv2(dec_out.transpose(2,1)).transpose(1,2)
          if self.output_attention:
              return dec_out[:,-self.pred_len:,:], attns
          else:
              return dec_out[:,-self.pred_len:,:] # [B, L, D]
  