LM has virtual 2nd batch

audiocraft/audiogen.py

@@ -11,9 +11,160 @@ and provide easy access to the generation API.
 
 import typing as tp
 import torch
-
-from audiocraft.genmodel import BaseGenModel
 from audiocraft.loaders import load_compression_model, load_lm_model
+import typing as tp
+import omegaconf
+import torch
+import numpy as np
+from abc import ABC, abstractmethod
+from .lm import LMModel
+from .conditioners import ConditioningAttributes
+from .utils.autocast import TorchAutocast
+
+
+
+def _shift(x):
+    n = x.shape[2]
+    i = np.random.randint(.24 * n, max(1, .74 * n))  # high should be above >= 0 TBD do we have very short segments
+    x = torch.roll(x, i, dims=2)
+    return x
+
+
+class BaseGenModel(ABC):
+    """Base generative model with convenient generation API.
+
+    Args:
+        name (str)
+        compression_model (CompressionModel): Encodec with Seanet Decoder
+        lm
+        max_duration (float, optional): As is using top250 token draw() we can gen xN sequences
+    """
+    def __init__(self,
+                 name, 
+                 compression_model,
+                 lm,
+                 max_duration=None):
+        self.name = name
+        self.compression_model = compression_model
+        self.lm = lm
+        self.cfg: tp.Optional[omegaconf.DictConfig] = None
+        # Just to be safe, let's put everything in eval mode.
+        self.compression_model.eval()
+        self.lm.eval()
+
+        if hasattr(lm, 'cfg'):
+            cfg = lm.cfg
+            assert isinstance(cfg, omegaconf.DictConfig)
+            self.cfg = cfg
+
+        if max_duration is None:
+            if self.cfg is not None:
+                max_duration = lm.cfg.dataset.segment_duration  # type: ignore
+            else:
+                raise ValueError("You must provide max_duration when building directly your GenModel")
+        assert max_duration is not None
+
+        self.max_duration: float = max_duration
+        self.duration = self.max_duration
+        self.device = next(iter(lm.parameters())).device
+        self.generation_params={}
+        
+        if self.device.type == 'cpu':
+            self.autocast = TorchAutocast(enabled=False)
+        else:
+            self.autocast = TorchAutocast(
+                enabled=True, 
+                device_type=self.device.type, 
+                dtype=torch.float16)
+
+    @property
+    def frame_rate(self) -> float:
+        """Roughly the number of AR steps per seconds."""
+        return self.compression_model.frame_rate
+
+    @property
+    def sample_rate(self) -> int:
+        """Sample rate of the generated audio."""
+        return self.compression_model.sample_rate
+
+    @property
+    def audio_channels(self) -> int:
+        """Audio channels of the generated audio."""
+        return self.compression_model.channels
+
+    @torch.no_grad()
+    def _prepare_tokens_and_attributes(
+            self,
+            descriptions,
+            prompt,
+    ):
+        attributes = [
+            ConditioningAttributes(text={'description': description}) for description in descriptions]
+        prompt_tokens = None
+        return attributes, prompt_tokens
+
+    def generate_unconditional(self, 
+                               num_samples,
+                               progress=False,
+                               return_tokens=False):
+        descriptions: tp.List[tp.Optional[str]] = [None] * num_samples
+        attributes, _ = self._prepare_tokens_and_attributes(descriptions, None)
+        tokens = self._generate_tokens(attributes)
+        if return_tokens:
+            return self.generate_audio(tokens), tokens
+        return self.generate_audio(tokens)
+
+    def generate(self, 
+                 descriptions, 
+                 progress=False, 
+                 return_tokens=False):
+        attributes, _ = self._prepare_tokens_and_attributes(descriptions, None)
+        tokens = self._generate_tokens(attributes)
+        if return_tokens:
+            return self.generate_audio(tokens), tokens
+        return self.generate_audio(tokens)
+
+    def _generate_tokens(self, attributes,
+                         prompt_tokens=None,
+                         progress=False):
+        
+        total_gen_len = int(self.duration * self.frame_rate)
+        max_prompt_len = int(min(self.duration, self.max_duration) * self.frame_rate)
+        current_gen_offset: int = 0
+
+        
+
+        
+
+        if self.duration <= self.max_duration:
+            # generate by sampling from LM, simple case.
+            
+            with self.autocast:
+                gen_tokens = self.lm.generate(conditions=attributes,
+                                                callback=None, 
+                                                max_gen_len=total_gen_len, 
+                                                **self.generation_params)
+        else:
+            print('<>Long gen ?<>')
+        # print(f'{gen_tokens.shape=}')   # [5,4,35]
+        # FLATTEN BATCH AS EXTRA SEQUENCE (BATCH IS VIRTUAL JUST MULTINOMIAL SAMPLING OF N_DRAW TOKENS)
+        gen_tokens = gen_tokens.transpose(0, 1).reshape(4, -1)[None, :, :]
+        for _ in range(3):
+            print(gen_tokens.shape)
+            gen_tokens = _shift(gen_tokens)
+        return gen_tokens
+
+    def generate_audio(self, gen_tokens: torch.Tensor) -> torch.Tensor:
+        """Generate Audio from tokens."""
+        assert gen_tokens.dim() == 3
+        with torch.no_grad():
+            gen_audio = self.compression_model.decode(gen_tokens, None)
+        return gen_audio
+
+
+
+
+
 
 class AudioGen(BaseGenModel):
     

audiocraft/conditioners.py

@@ -385,19 +385,9 @@ class ConditionFuser(StreamingModule):
 
     def forward(
         self,
-        input: torch.Tensor,
-        conditions: tp.Dict[str, ConditionType]
-    ) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
-        """Fuse the conditions to the provided model input.
-
-        Args:
-            input (torch.Tensor): Transformer input.
-            conditions (dict[str, ConditionType]): Dict of conditions.
-        Returns:
-            tuple[torch.Tensor, torch.Tensor]: The first tensor is the transformer input
-                after the conditions have been fused. The second output tensor is the tensor
-                used for cross-attention or None if no cross attention inputs exist.
-        """
+        input,
+        conditions):
+        
         B, T, _ = input.shape
 
 

audiocraft/genmodel.py

@@ -1,144 +0,0 @@
-import typing as tp
-import omegaconf
-import torch
-
-from abc import ABC, abstractmethod
-
-from .lm import LMModel
-from .conditioners import ConditioningAttributes
-from .utils.autocast import TorchAutocast
-
-
-class BaseGenModel(ABC):
-    """Base generative model with convenient generation API.
-
-    Args:
-        name (str): name of the model.
-        compression_model (CompressionModel): Encodec with Seanet Decoder
-        lm (LMModel): Language model over discrete representations
-        max_duration (float, optional): As is using top250 token draw() we can gen xN sequences
-    """
-    def __init__(self, name: str, compression_model, lm: LMModel,
-                 max_duration: tp.Optional[float] = None):
-        self.name = name
-        self.compression_model = compression_model
-        self.lm = lm
-        self.cfg: tp.Optional[omegaconf.DictConfig] = None
-        # Just to be safe, let's put everything in eval mode.
-        self.compression_model.eval()
-        self.lm.eval()
-
-        if hasattr(lm, 'cfg'):
-            cfg = lm.cfg
-            assert isinstance(cfg, omegaconf.DictConfig)
-            self.cfg = cfg
-
-        if max_duration is None:
-            if self.cfg is not None:
-                max_duration = lm.cfg.dataset.segment_duration  # type: ignore
-            else:
-                raise ValueError("You must provide max_duration when building directly your GenModel")
-        assert max_duration is not None
-
-        self.max_duration: float = max_duration
-        self.duration = self.max_duration
-        self.device = next(iter(lm.parameters())).device
-        self.generation_params={}
-        
-        if self.device.type == 'cpu':
-            self.autocast = TorchAutocast(enabled=False)
-        else:
-            self.autocast = TorchAutocast(
-                enabled=True, 
-                device_type=self.device.type, 
-                dtype=torch.float16)
-
-    @property
-    def frame_rate(self) -> float:
-        """Roughly the number of AR steps per seconds."""
-        return self.compression_model.frame_rate
-
-    @property
-    def sample_rate(self) -> int:
-        """Sample rate of the generated audio."""
-        return self.compression_model.sample_rate
-
-    @property
-    def audio_channels(self) -> int:
-        """Audio channels of the generated audio."""
-        return self.compression_model.channels
-
-    
-
-    @abstractmethod
-    def set_generation_params(self, *args, **kwargs):
-        """Set the generation parameters."""
-        raise NotImplementedError("No base implementation for setting generation params.")
-
-    @staticmethod
-    @abstractmethod
-    def get_pretrained(name: str, device=None):
-        raise NotImplementedError("No base implementation for getting pretrained model")
-
-    @torch.no_grad()
-    def _prepare_tokens_and_attributes(
-            self,
-            descriptions,
-            prompt,
-    ):
-        attributes = [
-            ConditioningAttributes(text={'description': description}) for description in descriptions]
-        prompt_tokens = None
-        return attributes, prompt_tokens
-
-    def generate_unconditional(self, 
-                               num_samples,
-                               progress=False,
-                               return_tokens=False):
-        descriptions: tp.List[tp.Optional[str]] = [None] * num_samples
-        attributes, _ = self._prepare_tokens_and_attributes(descriptions, None)
-        tokens = self._generate_tokens(attributes)
-        if return_tokens:
-            return self.generate_audio(tokens), tokens
-        return self.generate_audio(tokens)
-
-    def generate(self, descriptions, progress = False, return_tokens= False):
-        attributes, _ = self._prepare_tokens_and_attributes(descriptions, None)
-        tokens = self._generate_tokens(attributes)
-        if return_tokens:
-            return self.generate_audio(tokens), tokens
-        return self.generate_audio(tokens)
-
-    def _generate_tokens(self, attributes,
-                         prompt_tokens=None,
-                         progress=False):
-        
-        total_gen_len = int(self.duration * self.frame_rate)
-        max_prompt_len = int(min(self.duration, self.max_duration) * self.frame_rate)
-        current_gen_offset: int = 0
-
-        
-
-        
-
-        if self.duration <= self.max_duration:
-            # generate by sampling from LM, simple case.
-            
-            with self.autocast:
-                gen_tokens = self.lm.generate(conditions=attributes,
-                                                callback=None, 
-                                                max_gen_len=total_gen_len, 
-                                                **self.generation_params)
-        else:
-            print('<>Long gen ?<>')
-        # print(f'{gen_tokens.shape=}')   # [5,4,35]
-        # FLATTEN BATCH AS EXTRA SEQUENCE (BATCH IS VIRTUAL JUST MULTINOMIAL SAMPLING OF N_DRAW TOKENS)
-        gen_tokens = gen_tokens.transpose(0, 1).reshape(4, -1)[None, :, :]   
-        return gen_tokens
-
-    def generate_audio(self, gen_tokens: torch.Tensor) -> torch.Tensor:
-        """Generate Audio from tokens."""
-        assert gen_tokens.dim() == 3
-        with torch.no_grad():
-            gen_audio = self.compression_model.decode(gen_tokens, None)
-        return gen_audio

audiocraft/lm.py

@@ -148,7 +148,7 @@ class LMModel(StreamingModule):
         super().__init__()
         self.cfg_coef = cfg_coef
         
-        self.n_draw = 20
+        self.n_draw = 24
         self.condition_provider = condition_provider
         self.fuser = fuser
         self.card = card  # 2048 ?
@@ -213,24 +213,22 @@ class LMModel(StreamingModule):
     def num_codebooks(self) -> int:
         return self.n_q
 
-    def forward(self, 
+    def forward(self,
                 sequence,
                 conditions,
                 condition_tensors=None,
                 stage = -1):
         B, K, S = sequence.shape
-        assert K == self.num_codebooks, "Sequence shape must match the specified number of codebooks"
+        
+        
         input_ = sum([self.emb[k](sequence[:, k]) for k in range(K)])
-        if condition_tensors is None:
-            assert not self._is_streaming, "Conditions tensors should be precomputed when streaming."
-            
-            # encode conditions and fuse, both have a streaming cache to not recompute when generating.
-            condition_tensors = self.condition_provider(tokenized)
-        else:
-            assert not conditions, "Shouldn't pass both conditions and condition_tensors."
+        
         
         input_, cross_attention_input = self.fuser(input_, condition_tensors) # DEFINE conditioners.py
-
+        
+        # print(f'{input_.shape=}  {cross_attention_input.shape=}  FUSER LLM FORw')
+        # input_.shape=torch.Size([1, 1, 1536])  cross_attention_input.shape=torch.Size([2, 7, 1536])  FUSER LLM FORw
+        
         out = self.transformer(input_, cross_attention_src=cross_attention_input,
                                src_mask=(self.attn_mask_per_stage[stage] if stage >= 0 else None))
         if self.out_norm:
@@ -240,7 +238,7 @@ class LMModel(StreamingModule):
         # remove the prefix from the model outputs
         if len(self.fuser.fuse2cond['prepend']) > 0:
             logits = logits[:, :, -S:]
-            print('PRESFIX')
+            print('==========================================PRESFIX')
 
         return logits  # [B, K, S, card]
 
@@ -260,49 +258,20 @@ class LMModel(StreamingModule):
         cfg_coef = self.cfg_coef if cfg_coef is None else cfg_coef
         model = self if self._fsdp is None else self._fsdp
         two_step_cfg = self.two_step_cfg if two_step_cfg is None else two_step_cfg
-        if two_step_cfg and cfg_conditions != {}:
-            print('\nNOT HERE\n')
-        else:
-            print('C')
-            assert isinstance(cfg_conditions, dict)
-            condition_tensors = cfg_conditions
-            if condition_tensors:
-                print('\nDcat\n')  # enters here
-                
-                sequence = torch.cat([sequence, sequence], dim=0) # if i concatenate
-                # concatenates in batch but we only want to run 1st sequence - continutation
-                # the other paths will build "BLindly"
-            all_logits = model(
-                sequence,
-                conditions=[], condition_tensors=condition_tensors)
-            if condition_tensors:
-                cond_logits, uncond_logits = all_logits.split(B, dim=0) #torch.Size([2, 4, 1, 2048])
-                # logits = uncond_logits + (cond_logits - uncond_logits) * cfg_coef
-                # logits = 3 * cond_logits - 2.4 * uncond_logits
-                logits = 2 * cond_logits - 1.4 * uncond_logits
-            else:
-                print('\nF!\n')
-                
-
-        logits = logits.permute(0, 1, 3, 2)  # [1, 4, 2048, 1]
-        # No crop this is just squeeze() of time
-        logits = logits[..., -1]  # [1 x 4 x 2048]
-        
-
-        # Apply softmax for sampling if temp > 0. Else, do greedy sampling to avoid zero division error.
+        condition_tensors = cfg_conditions
         
-        # print(f'\nR {temp=} {top_p=} {top_k=}\n') -------------> R temp=1.0 top_p=0.0 top_k=250
-        # print(f'{temp=}')  # 1.0
-        probs = torch.softmax(logits / temp, dim=-1)
         
-        next_token = utils.sample_top_k(probs, k=top_k, n_draw=self.n_draw)
         
+        logits = model(
+            sequence,                                          # cond_logits = wav condition 
+            conditions=[], condition_tensors=condition_tensors)  # uncond_logits already see the text
         
-        # th decoder will smooth the transitions
-        # so if we have 2 tokens although the 2nd token we need it for replica later
-        # so let it as batch and reshape at the final time-inversion
+        # print(f'{logits.shape=} L')
+        logits = logits[0, :, :, :].transpose(1,0)  # sample expects [1, 4, 2048]
+        # logits = [2, 4, 1, 2048]
+        # print(f'{B=}, {logits.shape=} SAMPLER {top_k=}')
+        next_token = utils.sample_top_k(logits, k=top_k, n_draw=self.n_draw)  # [1,4,2048] logits
         
-        # To return multiple tokens here (batch_size = num_draws)
         return next_token
 
     # GENERATE class revert_codebook_patterns()
@@ -385,7 +354,7 @@ class LMModel(StreamingModule):
             #   but continue the sequence only with isingle next token
             
             for offset in range(1, gen_sequence_len):  # start_offset_sequence=1
-                print(f'{offset=}')
+                # print(f'{_gen_sequence.shape=}')  # [1,4,16]
                 # starts from 1 not 0 thus uses the 0:1 as curr sequence
                 # although this is empty contains -1 ?
     

audiocraft/transformer.py

@@ -177,7 +177,7 @@ class StreamingMultiheadAttention(StreamingModule):
         self.past_context = past_context
         self.memory_efficient = memory_efficient
         self.attention_as_float32 = attention_as_float32
-        self.rope = rope
+        
         self.cross_attention = cross_attention
         self.safe_streaming = safe_streaming
         self.num_heads = num_heads
@@ -230,38 +230,6 @@ class StreamingMultiheadAttention(StreamingModule):
                     state_dict[prefix + "mha." + key] = state_dict.pop(prefix + key)
         super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
 
-    def _get_mask(self, current_steps: int, device: torch.device, dtype: torch.dtype):
-        # Return a causal mask, accounting for potentially stored past keys/values
-        # We actually return a bias for the attention score, as this has the same
-        # convention both in the builtin MHA in Pytorch, and Xformers functions.
-        time_dim = _get_attention_time_dimension(self.memory_efficient)
-        if self.memory_efficient:
-            from xformers.ops import LowerTriangularMask
-            if current_steps == 1:
-                # If we only have one step, then we do not need a mask.
-                return None
-            elif 'past_keys' in self._streaming_state:
-                raise RuntimeError("Not supported at the moment")
-            else:
-                # Then we can safely use a lower triangular mask
-                return LowerTriangularMask()
-        if self._streaming_state:
-            past_keys = self._streaming_state['past_keys']
-            past_steps = past_keys.shape[time_dim]
-        else:
-            past_steps = 0
-
-        queries_pos = torch.arange(
-            past_steps, current_steps + past_steps, device=device).view(-1, 1)
-        keys_pos = torch.arange(past_steps + current_steps, device=device).view(1, -1)
-        delta = queries_pos - keys_pos
-        valid = delta >= 0
-        if self.past_context is not None:
-            valid &= (delta <= self.past_context)
-        return torch.where(
-            valid,
-            torch.zeros([], device=device, dtype=dtype),
-            torch.full([], float('-inf'), device=device, dtype=dtype))
 
     def _complete_kv(self, k, v):
         time_dim = _get_attention_time_dimension(self.memory_efficient)
@@ -272,11 +240,14 @@ class StreamingMultiheadAttention(StreamingModule):
             return k, v
         # Complete the key/value pair using the streaming state.
         if self._streaming_state:
+            # print('{self._streaming_state.keys()=}')   EMPTY - ALTHOUGH WE HAVE STREAMING STATE
             pk = self._streaming_state['past_keys']
             nk = torch.cat([pk, k], dim=time_dim)
             if v is k:
+                
                 nv = nk
             else:
+                
                 pv = self._streaming_state['past_values']
                 nv = torch.cat([pv, v], dim=time_dim)
         else:
@@ -286,35 +257,28 @@ class StreamingMultiheadAttention(StreamingModule):
         assert nk.shape[time_dim] == nv.shape[time_dim]
         offset = 0
         if self.past_context is not None:
+            
             offset = max(0, nk.shape[time_dim] - self.past_context)
         if self._is_streaming:
             self._streaming_state['past_keys'] = nk[:, offset:]
             if v is not k:
+                
                 self._streaming_state['past_values'] = nv[:, offset:]
             if 'offset' in self._streaming_state:
+                
                 self._streaming_state['offset'] += offset
             else:
+            
                 self._streaming_state['offset'] = torch.tensor(0)
         return nk, nv
 
-    def _apply_rope(self, query: torch.Tensor, key: torch.Tensor):
-        time_dim = _get_attention_time_dimension(self.memory_efficient)
-        # Apply rope embeddings to query and key tensors.
-        assert self.rope is not None
-        if 'past_keys' in self._streaming_state:
-            past_keys_offset = self._streaming_state['past_keys'].shape[1]
-        else:
-            past_keys_offset = 0
-        if 'offset' in self._streaming_state:
-            past_context_offset = int(self._streaming_state['offset'].item())
-        else:
-            past_context_offset = 0
-        streaming_offset = past_context_offset + past_keys_offset
-        return self.rope.rotate_qk(query, key, start=streaming_offset, time_dim=time_dim)
+    
 
     def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
                 key_padding_mask=None, need_weights=False, attn_mask=None,
                 average_attn_weights=True, is_causal=False):
+        
+        
         assert not is_causal, ("New param added in torch 2.0.1 not supported, "
                                "use the causal args in the constructor.")
 
@@ -328,29 +292,22 @@ class StreamingMultiheadAttention(StreamingModule):
             assert self.causal or self.cross_attention, \
                 "Streaming only available for causal or cross attention"
 
-        custom_attn_mask = attn_mask is not None
+        
 
-        if self.causal:
-            assert attn_mask is None
-            # At the moment we specialize only for the self-attention case.
-            assert query.shape[1] == key.shape[1], "Causal only for same length query / key / value"
-            assert value.shape[1] == key.shape[1], "Causal only for same length query / key / value"
-            attn_mask = self._get_mask(query.shape[1], query.device, query.dtype)
+        
+            
+            
 
         if self.custom:
-            # custom implementation
-            assert need_weights is False
-            assert key_padding_mask is None
+            
             if self.cross_attention:
+                
                 # Different queries, keys, values, we have to spit manually the weights
                 # before applying the linear.
                 dim = self.in_proj_weight.shape[0] // 3
-                if self.in_proj_bias is None:
-                    bias_q, bias_k, bias_v = None, None, None
-                else:
-                    bias_q = self.in_proj_bias[:dim]
-                    bias_k = self.in_proj_bias[dim: 2 * dim]
-                    bias_v = self.in_proj_bias[2 * dim:]
+                
+                bias_q, bias_k, bias_v = None, None, None
+                
                 q = nn.functional.linear(query, self.in_proj_weight[:dim], bias_q)
                 # todo: when streaming, we could actually save k, v and check the shape actually match.
                 k = nn.functional.linear(key, self.in_proj_weight[dim: 2 * dim], bias_k)
@@ -366,125 +323,31 @@ class StreamingMultiheadAttention(StreamingModule):
                     assert value is key, "specialized implementation"
                 projected = nn.functional.linear(query, self.in_proj_weight, self.in_proj_bias)
                 if self.kv_repeat == 1:
+                    
                     if time_dim == 2:
                         bound_layout = "b h p t d"
                     else:
                         bound_layout = "b t p h d"
                     packed = rearrange(projected, f"b t (p h d) -> {bound_layout}", p=3, h=self.num_heads)
                     q, k, v = ops.unbind(packed, dim=2)
-                else:
-                    embed_dim = self.embed_dim
-                    per_head_dim = (embed_dim // self.num_heads)
-                    kv_heads = self.num_heads // self.kv_repeat
-                    q = projected[:, :, :embed_dim]
-                    start = embed_dim
-                    end = start + per_head_dim * kv_heads
-                    k = projected[:, :, start: end]
-                    v = projected[:, :, end:]
-                    q = rearrange(q, f"b t (h d) -> {layout}", h=self.num_heads)
-                    k = rearrange(k, f"b t (h d) -> {layout}", h=kv_heads)
-                    v = rearrange(v, f"b t (h d) -> {layout}", h=kv_heads)
 
-                if self.qk_layer_norm is True:
-                    assert self.kv_repeat == 1
-                    q, k = [rearrange(x, f"{layout} -> b t (h d)") for x in [q, k]]
-                    q = self.q_layer_norm(q)
-                    k = self.k_layer_norm(k)
-                    q, k = [rearrange(x, f"b t (h d) -> {layout}", h=self.num_heads) for x in [q, k]]
-                if self.rope:
-                    q, k = self._apply_rope(q, k)
+                
                 k, v = self._complete_kv(k, v)
-                if self.kv_repeat > 1:
-                    k = expand_repeated_kv(k, self.kv_repeat, self.memory_efficient)
-                    v = expand_repeated_kv(v, self.kv_repeat, self.memory_efficient)
-            if self.attention_as_float32:
-                q, k, v = [x.float() for x in [q, k, v]]
-            if self.memory_efficient:
-                if custom_attn_mask:
-                    # When using a custom attn mask:
-                    # Move to query's device, repeat for each sample, remove align8 padding
-                    seq_len = query.shape[1]
-                    attn_mask = attn_mask.to(q.dtype)
-                    attn_mask = attn_mask.repeat((q.shape[0], 1, 1, 1))
-                    attn_mask = attn_mask[..., :seq_len, :seq_len]
-
-                p = self.dropout if self.training else 0
-                if _efficient_attention_backend == 'torch':
-                    x = torch.nn.functional.scaled_dot_product_attention(
-                        q, k, v, is_causal=attn_mask is not None, dropout_p=p)
-                else:
-                    x = ops.memory_efficient_attention(q, k, v, attn_mask, p=p)
-            else:
-                # We include the dot product as float32, for consistency
-                # with the other implementations that include that step
-                # as part of the attention. Note that when using `autocast`,
-                # the einsums would be done as bfloat16, but the softmax
-                # would be done as bfloat16, so `attention_as_float32` will
-                # extend a bit the range of operations done in float32,
-                # although this should make no difference.
-                q = q / q.shape[-1] ** 0.5
-                key_layout = layout.replace('t', 'k')
-                query_layout = layout
-                if self._is_streaming and self.safe_streaming and q.device.type == 'cuda':
-                    with torch.autocast(device_type=q.device.type, dtype=torch.float32):
-                        pre_w = torch.einsum(f"{query_layout},{key_layout}-> b h t k", q, k)
-                else:
-                    pre_w = torch.einsum(f"{query_layout},{key_layout}-> b h t k", q, k)
-                if attn_mask is not None:
-                    pre_w = pre_w + attn_mask
-                w = torch.softmax(pre_w, dim=-1)
-                w = F.dropout(w, self.dropout, training=self.training).to(v)
-                # Key and value have the same format.
-                x = torch.einsum(f"b h t k, {key_layout} -> {layout}", w, v)
+#print(f'{k.shape=}, {v.shape=}, {q.shape=}\n\n\n\n')
+            # what is the 24 dimension is this heads?
+
+            x = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, is_causal=attn_mask is not None, dropout_p=0)
+                
+            
             x = x.to(dtype)
             x = rearrange(x, f"{layout} -> b t (h d)", h=self.num_heads)
             x = self.out_proj(x)
-        else:
-            key, value = self._complete_kv(key, value)
-            if self.attention_as_float32:
-                query, key, value = [x.float() for x in [query, key, value]]
-            x, _ = self.mha(
-                query, key, value, key_padding_mask,
-                need_weights, attn_mask, average_attn_weights)
-            x = x.to(dtype)
-
         return x, None
 
 
 class StreamingTransformerLayer(nn.TransformerEncoderLayer):
-    """TransformerLayer with Streaming / Causal support.
-    This also integrates cross_attention, when passing `cross_attention=True`,
-    rather than having two separate classes like in PyTorch.
 
-    Args:
-        d_model (int): Dimension of the data.
-        num_heads (int): Number of heads.
-        dim_feedforward (int): Intermediate dimension of FF module.
-        dropout (float): Dropout both for MHA and FF.
-        bias_ff (bool): Use bias for FF.
-        bias_attn (bool): Use bias for MHA.
-        causal (bool): Causal mask applied automatically.
-        past_context (int, optional): Receptive field for the causal mask, infinite if None.
-        custom (bool): Use custom MHA implementation, for testing / benchmarking.
-        memory_efficient (bool): Use xformers based memory efficient attention.
-        attention_as_float32 (bool): Perform the attention as float32
-            (especially important with memory_efficient as autocast won't do this automatically).
-        qk_layer_norm (bool): Layer normalization applied to queries and keys before dot product in attention.
-        qk_layer_norm_cross (bool): Same for the cross attention.
-        cross_attention (bool): If True, expect to get secondary input for cross-attention.
-            Cross attention will use the default MHA, as it typically won't require
-            special treatment.
-        layer_scale (float, optional): If not None, LayerScale will be used with
-            the given value as initial scale.
-        rope (`RotaryEmbedding`, optional): Rope embedding to use.
-        attention_dropout (float, optional): If not None, separate the value of the dimension dropout
-            in FFN and of the attention dropout.
-        kv_repeat (int): If > 1, will repeat keys and queries multiple times (need to divide num_heads).
-            This will lead to faster decoding time on A100 or other GPUs with tensorcore.
-        device (torch.device, optional): Device on which to initialize.
-        dtype (torch.dtype, optional): dtype to use.
-        **kwargs: See `nn.TransformerEncoderLayer`.
-    """
     def __init__(self, d_model: int, num_heads: int, dim_feedforward: int = 2048, dropout: float = 0.1,
                  bias_ff: bool = True, bias_attn: bool = True, causal: bool = False,
                  past_context: tp.Optional[int] = None, custom: bool = False,
@@ -632,6 +495,7 @@ class StreamingTransformer(StreamingModule):
         assert positional_embedding in ['sin', 'rope', 'sin_rope']
         self.rope: tp.Optional[RotaryEmbedding] = None
         if self.positional_embedding in ['rope', 'sin_rope']:
+            print('ROPE\nL')
             assert _is_custom(custom, memory_efficient)
             self.rope = RotaryEmbedding(d_model // num_heads, max_period=max_period,
                                         xpos=xpos, scale=positional_scale, device=device)
@@ -659,43 +523,11 @@ class StreamingTransformer(StreamingModule):
                 # backward hook inside of FSDP...
                 layer._magma_checkpointed = True  # type: ignore
 
-    def _apply_layer(self, layer, *args, **kwargs):
-        method = self.checkpointing
-        print(f'{method=}')
-        if method == 'none':
-            print([i.shape for i in args])
-            x = layer(*args, **kwargs)  # [10, 1, 1536]   probably does no t detect the bathc somwhere
-            return x
-        # elif method == 'torch':
-        #     print('TORCH')
-        #     return torch_checkpoint(layer, *args, use_reentrant=False, **kwargs)
-        # elif method.startswith('xformers'):
-        #     print('XFORMERS')
-        #     from xformers.checkpoint_fairinternal import checkpoint, _get_default_policy
-        #     if method == 'xformers_default':
-        #         # those operations will be saved, and not recomputed.
-        #         # According to Francisco we can get smarter policies but this is a good start.
-        #         allow_list = [
-        #             "xformers.efficient_attention_forward_cutlass.default",
-        #             "xformers_flash.flash_fwd.default",
-        #             "aten.addmm.default",
-        #             "aten.mm.default",
-        #         ]
-        #     elif method == 'xformers_mm':
-        #         # those operations will be saved, and not recomputed.
-        #         # According to Francisco we can get smarter policies but this is a good start.
-        #         allow_list = [
-        #             "aten.addmm.default",
-        #             "aten.mm.default",
-        #         ]
-        #     else:
-        #         raise ValueError(f"xformers checkpointing xformers policy {method} is not known.")
-        #     policy_fn = _get_default_policy(allow_list)
-        #     return checkpoint(layer, *args, policy_fn=policy_fn, **kwargs)
-        # else:
-        #     raise ValueError(f"Checkpointing method {method} is unknown.")
+    
 
     def forward(self, x: torch.Tensor, *args, **kwargs):
+        # Input  x: [1, 1, 1536]
+        # Output x: [2, 1, 1536]  how is batch expanded to 2
         B, T, C = x.shape
 
         if 'offsets' in self._streaming_state:
@@ -704,17 +536,20 @@ class StreamingTransformer(StreamingModule):
             offsets = torch.zeros(B, dtype=torch.long, device=x.device)
 
         if self.positional_embedding in ['sin', 'sin_rope']:
+            # print(f'{self.positional_embedding=}\n')  'sin'
             positions = torch.arange(T, device=x.device).view(1, -1, 1)
             positions = positions + offsets.view(-1, 1, 1)
             pos_emb = create_sin_embedding(positions, C, max_period=self.max_period, dtype=x.dtype)
             x = x + self.positional_scale * pos_emb
 
         for layer in self.layers:
-            x = self._apply_layer(layer, x, *args, **kwargs)
+            # print(f'{args=} {kwargs.keys()=}')   
+            # # kwargs=() kwargs={'cross_attention_src', 'src_mask'}
+            x = layer(x, **kwargs)
 
         if self._is_streaming:
             self._streaming_state['offsets'] = offsets + T
-
+        print('OUT STReamTransfor', x.shape)
         return x
 
     def make_optim_group(self):

audiocraft/utils/utils.py

@@ -94,6 +94,11 @@ def sample_top_k(p, k, n_draw=None):
         p probabs 2048 ?
         num_draw : how many tokens to sample (for duplicate elongation)
     """
+
+    p = torch.softmax(p / 1.0, dim=-1)
+         
+
+    
     top_k_value, i250 = torch.topk(p, k, dim=-1)   # probs: [1, 4, 2048]
     min_value_top_k = top_k_value[..., [-1]]  # 
     p *= (p >= min_value_top_k).float()