factor diffusion

Modules/diffusion/diffusion.py

@@ -1,85 +0,0 @@
-from math import pi
-from random import randint
-from typing import Any, Optional, Sequence, Tuple, Union
-
-import torch
-from einops import rearrange
-from torch import Tensor, nn
-from tqdm import tqdm
-
-from .utils import *
-from .sampler import *
-
-"""
-Diffusion Classes (generic for 1d data)
-"""
-
-
-class Model1d(nn.Module):
-    def __init__(self, unet_type: str = "base", **kwargs):
-        super().__init__()
-        diffusion_kwargs, kwargs = groupby("diffusion_", kwargs)
-        self.unet = None
-        self.diffusion = None
-
-    def forward(self, x: Tensor, **kwargs) -> Tensor:
-        return self.diffusion(x, **kwargs)
-
-    def sample(self, *args, **kwargs) -> Tensor:
-        return self.diffusion.sample(*args, **kwargs)
-
-
-"""
-Audio Diffusion Classes (specific for 1d audio data)
-"""
-
-
-def get_default_model_kwargs():
-    return dict(
-        channels=128,
-        patch_size=16,
-        multipliers=[1, 2, 4, 4, 4, 4, 4],
-        factors=[4, 4, 4, 2, 2, 2],
-        num_blocks=[2, 2, 2, 2, 2, 2],
-        attentions=[0, 0, 0, 1, 1, 1, 1],
-        attention_heads=8,
-        attention_features=64,
-        attention_multiplier=2,
-        attention_use_rel_pos=False,
-        diffusion_type="v",
-        diffusion_sigma_distribution=UniformDistribution(),
-    )
-
-
-def get_default_sampling_kwargs():
-    return dict(sigma_schedule=LinearSchedule(), sampler=VSampler(), clamp=True)
-
-class AudioDiffusionConditional(Model1d):
-    def __init__(
-        self,
-        embedding_features: int,
-        embedding_max_length: int,
-        embedding_mask_proba: float = 0.1,
-        **kwargs,
-    ):
-        self.embedding_mask_proba = embedding_mask_proba
-        default_kwargs = dict(
-            **get_default_model_kwargs(),
-            unet_type="cfg",
-            context_embedding_features=embedding_features,
-            context_embedding_max_length=embedding_max_length,
-        )
-        super().__init__(**{**default_kwargs, **kwargs})
-
-    def forward(self, *args, **kwargs):
-        default_kwargs = dict(embedding_mask_proba=self.embedding_mask_proba)
-        return super().forward(*args, **{**default_kwargs, **kwargs})
-
-    def sample(self, *args, **kwargs):
-        default_kwargs = dict(
-            **get_default_sampling_kwargs(),
-            embedding_scale=5.0,
-        )
-        return super().sample(*args, **{**default_kwargs, **kwargs})
-
-

Modules/diffusion/modules.py

@@ -1,8 +1,5 @@
 from math import floor, log, pi
-from typing import Any, List, Optional, Sequence, Tuple, Union
-
-from .utils import *
-
+import torch.nn.functional as F
 import torch
 import torch.nn as nn
 from einops import rearrange, reduce, repeat
@@ -11,9 +8,10 @@ from einops_exts import rearrange_many
 from torch import Tensor, einsum
 
 
-"""
-Utils
-"""
+def default(val, d):
+    if val is not None: #exists(val):
+        return val
+    return d # d() if isfunction(d) else d
 
 class AdaLayerNorm(nn.Module):
     def __init__(self, style_dim, channels, eps=1e-5):
@@ -38,6 +36,9 @@ class AdaLayerNorm(nn.Module):
         return x.transpose(1, -1).transpose(-1, -2)
 
 class StyleTransformer1d(nn.Module):
+    
+    # artificial_stylets / models.py
+    
     def __init__(
         self,
         num_layers: int,
@@ -48,14 +49,14 @@ class StyleTransformer1d(nn.Module):
         use_context_time: bool = True,
         use_rel_pos: bool = False,
         context_features_multiplier: int = 1,
-        rel_pos_num_buckets: Optional[int] = None,
-        rel_pos_max_distance: Optional[int] = None,
-        context_features: Optional[int] = None,
-        context_embedding_features: Optional[int] = None,
-        embedding_max_length: int = 512,
+        # rel_pos_num_buckets: Optional[int] = None,
+        # rel_pos_max_distance: Optional[int] = None,
+        context_features=None,
+        context_embedding_features=None,
+        embedding_max_length=512,
     ):
         super().__init__()
-
+        
         self.blocks = nn.ModuleList(
             [
                 StyleTransformerBlock(
@@ -65,8 +66,8 @@ class StyleTransformer1d(nn.Module):
                     multiplier=multiplier,
                     style_dim=context_features,
                     use_rel_pos=use_rel_pos,
-                    rel_pos_num_buckets=rel_pos_num_buckets,
-                    rel_pos_max_distance=rel_pos_max_distance,
+                    # rel_pos_num_buckets=rel_pos_num_buckets,
+                    # rel_pos_max_distance=rel_pos_max_distance,
                 )
                 for i in range(num_layers)
             ]
@@ -81,11 +82,14 @@ class StyleTransformer1d(nn.Module):
             ),
         )
         
-        use_context_features = exists(context_features)
+        use_context_features = context_features is not None
         self.use_context_features = use_context_features
         self.use_context_time = use_context_time
 
         if use_context_time or use_context_features:
+            # print(f'{use_context_time=} {use_context_features=}ooooooooooooooooooooooooooooooooooo')
+            # raise ValueError
+            # True True  both context
             context_mapping_features = channels + context_embedding_features
 
             self.to_mapping = nn.Sequential(
@@ -96,7 +100,7 @@ class StyleTransformer1d(nn.Module):
             )
         
         if use_context_time:
-            assert exists(context_mapping_features)
+            
             self.to_time = nn.Sequential(
                 TimePositionalEmbedding(
                     dim=channels, out_features=context_mapping_features
@@ -105,7 +109,7 @@ class StyleTransformer1d(nn.Module):
             )
 
         if use_context_features:
-            assert exists(context_features) and exists(context_mapping_features)
+            
             self.to_features = nn.Sequential(
                 nn.Linear(
                     in_features=context_features, out_features=context_mapping_features
@@ -119,23 +123,23 @@ class StyleTransformer1d(nn.Module):
         
 
     def get_mapping(
-        self, time: Optional[Tensor] = None, features: Optional[Tensor] = None
-    ) -> Optional[Tensor]:
+        self, 
+        time=None,
+        features=None):
         """Combines context time features and features into mapping"""
         items, mapping = [], None
         # Compute time features
         if self.use_context_time:
-            assert_message = "use_context_time=True but no time features provided"
-            assert exists(time), assert_message
+
             items += [self.to_time(time)]
         # Compute features
         if self.use_context_features:
-            assert_message = "context_features exists but no features provided"
-            assert exists(features), assert_message
+
             items += [self.to_features(features)]
 
         # Compute joint mapping
         if self.use_context_time or self.use_context_features:
+            # raise ValueError
             mapping = reduce(torch.stack(items), "n b m -> b m", "sum")
             mapping = self.to_mapping(mapping)
 
@@ -160,8 +164,8 @@ class StyleTransformer1d(nn.Module):
     def forward(self, x: Tensor, 
                 time: Tensor, 
                 embedding_mask_proba: float = 0.0,
-                embedding: Optional[Tensor] = None, 
-                features: Optional[Tensor] = None,
+                embedding= None, 
+                features = None,
                embedding_scale: float = 1.0) -> Tensor:
         
         b, device = embedding.shape[0], embedding.device
@@ -174,13 +178,18 @@ class StyleTransformer1d(nn.Module):
             embedding = torch.where(batch_mask, fixed_embedding, embedding)
 
         if embedding_scale != 1.0:
-            # Compute both normal and fixed embedding outputs
+            
+            
             out = self.run(x, time, embedding=embedding, features=features)
             out_masked = self.run(x, time, embedding=fixed_embedding, features=features)
-            # Scale conditional output using classifier-free guidance
+            
+            raise ValueError
             return out_masked + (out - out_masked) * embedding_scale
+            
         else:
+            # raise ValueError
             return self.run(x, time, embedding=embedding, features=features)
+          
         
         return x
 
@@ -194,42 +203,45 @@ class StyleTransformerBlock(nn.Module):
         style_dim: int,
         multiplier: int,
         use_rel_pos: bool,
-        rel_pos_num_buckets: Optional[int] = None,
-        rel_pos_max_distance: Optional[int] = None,
-        context_features: Optional[int] = None,
+        # rel_pos_num_buckets: Optional[int] = None,
+        # rel_pos_max_distance: Optional[int] = None,
+        context_features = None,
     ):
         super().__init__()
 
-        self.use_cross_attention = exists(context_features) and context_features > 0
-
+        self.use_cross_attention = (context_features is not None) and (context_features > 0)
+        # print(f'{rel_pos_num_buckets=} {rel_pos_max_distance=}')  # None None
+        # raise ValueError
         self.attention = StyleAttention(
             features=features,
             style_dim=style_dim,
             num_heads=num_heads,
             head_features=head_features,
             use_rel_pos=use_rel_pos,
-            rel_pos_num_buckets=rel_pos_num_buckets,
-            rel_pos_max_distance=rel_pos_max_distance,
+            # rel_pos_num_buckets=rel_pos_num_buckets,
+            # rel_pos_max_distance=rel_pos_max_distance,
         )
 
         if self.use_cross_attention:
-            self.cross_attention = StyleAttention(
-                features=features,
-                style_dim=style_dim,
-                num_heads=num_heads,
-                head_features=head_features,
-                context_features=context_features,
-                use_rel_pos=use_rel_pos,
-                rel_pos_num_buckets=rel_pos_num_buckets,
-                rel_pos_max_distance=rel_pos_max_distance,
-            )
+            raise ValueError
+            # self.cross_attention = StyleAttention(
+            #     features=features,
+            #     style_dim=style_dim,
+            #     num_heads=num_heads,
+            #     head_features=head_features,
+            #     context_features=context_features,
+            #     use_rel_pos=use_rel_pos,
+            #     rel_pos_num_buckets=rel_pos_num_buckets,
+            #     rel_pos_max_distance=rel_pos_max_distance,
+            # )
 
         self.feed_forward = FeedForward(features=features, multiplier=multiplier)
 
-    def forward(self, x: Tensor, s: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
+    def forward(self, x: Tensor, s: Tensor, *, context = None) -> Tensor:
         x = self.attention(x, s) + x
         if self.use_cross_attention:
-            x = self.cross_attention(x, s, context=context) + x
+            raise ValueError
+            # x = self.cross_attention(x, s, context=context) + x
         x = self.feed_forward(x) + x
         return x
 
@@ -241,10 +253,10 @@ class StyleAttention(nn.Module):
         style_dim: int,
         head_features: int,
         num_heads: int,
-        context_features: Optional[int] = None,
+        context_features = None,
         use_rel_pos: bool,
-        rel_pos_num_buckets: Optional[int] = None,
-        rel_pos_max_distance: Optional[int] = None,
+        # rel_pos_num_buckets: Optional[int] = None,
+        # rel_pos_max_distance: Optional[int] = None,
     ):
         super().__init__()
         self.context_features = context_features
@@ -264,15 +276,16 @@ class StyleAttention(nn.Module):
             num_heads=num_heads,
             head_features=head_features,
             use_rel_pos=use_rel_pos,
-            rel_pos_num_buckets=rel_pos_num_buckets,
-            rel_pos_max_distance=rel_pos_max_distance,
+            # rel_pos_num_buckets=rel_pos_num_buckets,
+            # rel_pos_max_distance=rel_pos_max_distance,
         )
 
-    def forward(self, x: Tensor, s: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
-        assert_message = "You must provide a context when using context_features"
-        assert not self.context_features or exists(context), assert_message
+    def forward(self, x: Tensor, s: Tensor, *, context = None):
+        
+        # raise ValueError
         # Use context if provided
         context = default(context, x)
+        # print(context.shape,'ppppppppppppppppppppppppppppppppppppppppppp')  # bs, time, 1024
         # Normalize then compute q from input and k,v from context
         x, context = self.norm(x, s), self.norm_context(context, s)
         
@@ -280,7 +293,9 @@ class StyleAttention(nn.Module):
         # Compute and return attention
         return self.attention(q, k, v)
 
-def FeedForward(features: int, multiplier: int) -> nn.Module:
+
+def FeedForward(features, 
+                multiplier):
     mid_features = features * multiplier
     return nn.Sequential(
         nn.Linear(in_features=features, out_features=mid_features),
@@ -292,14 +307,14 @@ def FeedForward(features: int, multiplier: int) -> nn.Module:
 class AttentionBase(nn.Module):
     def __init__(
         self,
-        features: int,
+        features,
         *,
-        head_features: int,
-        num_heads: int,
-        use_rel_pos: bool,
-        out_features: Optional[int] = None,
-        rel_pos_num_buckets: Optional[int] = None,
-        rel_pos_max_distance: Optional[int] = None,
+        head_features,
+        num_heads,
+        use_rel_pos,
+        out_features = None,
+        # rel_pos_num_buckets: Optional[int] = None,
+        # rel_pos_max_distance: Optional[int] = None,
     ):
         super().__init__()
         self.scale = head_features ** -0.5
@@ -320,7 +335,11 @@ class AttentionBase(nn.Module):
         q, k, v = rearrange_many((q, k, v), "b n (h d) -> b h n d", h=self.num_heads)
         # Compute similarity matrix
         sim = einsum("... n d, ... m d -> ... n m", q, k)
-        sim = (sim + self.rel_pos(*sim.shape[-2:])) if self.use_rel_pos else sim
+        
+        #                     _____THERE_IS_NO_rel_po
+        # sim = (sim + self.rel_pos(*sim.shape[-2:])) if self.use_rel_pos else sim
+        # print(self.rel_pos)
+        
         sim = sim * self.scale
         # Get attention matrix with softmax
         attn = sim.softmax(dim=-1)
@@ -333,15 +352,15 @@ class AttentionBase(nn.Module):
 class Attention(nn.Module):
     def __init__(
         self,
-        features: int,
+        features,
         *,
-        head_features: int,
-        num_heads: int,
-        out_features: Optional[int] = None,
-        context_features: Optional[int] = None,
-        use_rel_pos: bool,
-        rel_pos_num_buckets: Optional[int] = None,
-        rel_pos_max_distance: Optional[int] = None,
+        head_features,
+        num_heads,
+        out_features=None,
+        context_features=None,
+        use_rel_pos,
+        # rel_pos_num_buckets: Optional[int] = None,
+        # rel_pos_max_distance: Optional[int] = None,
     ):
         super().__init__()
         self.context_features = context_features
@@ -363,13 +382,13 @@ class Attention(nn.Module):
             num_heads=num_heads,
             head_features=head_features,
             use_rel_pos=use_rel_pos,
-            rel_pos_num_buckets=rel_pos_num_buckets,
-            rel_pos_max_distance=rel_pos_max_distance,
+            # rel_pos_num_buckets=rel_pos_num_buckets,
+            # rel_pos_max_distance=rel_pos_max_distance,
         )
 
-    def forward(self, x: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
-        assert_message = "You must provide a context when using context_features"
-        assert not self.context_features or exists(context), assert_message
+    def forward(self, x: Tensor, *, context = None) -> Tensor:
+        # assert_message = "You must provide a context when using context_features"
+        # assert not self.context_features or exists(context), assert_message
         # Use context if provided
         context = default(context, x)
         # Normalize then compute q from input and k,v from context

Modules/diffusion/sampler.py

@@ -1,11 +1,59 @@
 from math import atan, cos, pi, sin, sqrt
-from typing import Any, Callable, List, Optional, Tuple, Type
-import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from einops import rearrange
 from torch import Tensor
-from .utils import *
+
+from functools import reduce
+from inspect import isfunction
+from math import ceil, floor, log2, pi
+# from typing import Callable, Dict, List, Optional, Sequence, Tuple, TypeVar, Union
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Generator, Tensor
+
+
+
+
+
+
+# def is_sequence(obj: T) -> TypeGuard[Union[list, tuple]]:
+#     return isinstance(obj, list) or isinstance(obj, tuple)
+
+
+def default(val, d):
+    if val is not None: #exists(val):
+        return val
+    return d #d() if isfunction(d) else d
+
+
+# def to_list(val: Union[T, Sequence[T]]) -> List[T]:
+#     if isinstance(val, tuple):
+#         return list(val)
+#     if isinstance(val, list):
+#         return val
+#     return [val]  # type: ignore
+
+
+# def prod(vals: Sequence[int]) -> int:
+#     return reduce(lambda x, y: x * y, vals)
+
+
+def closest_power_2(x: float) -> int:
+    exponent = log2(x)
+    distance_fn = lambda z: abs(x - 2 ** z)  # noqa
+    exponent_closest = min((floor(exponent), ceil(exponent)), key=distance_fn)
+    return 2 ** int(exponent_closest)
+
+def rand_bool(shape, proba, device = None):
+    if proba == 1:
+        return torch.ones(shape, device=device, dtype=torch.bool)
+    elif proba == 0:
+        return torch.zeros(shape, device=device, dtype=torch.bool)
+    else:
+        return torch.bernoulli(torch.full(shape, proba, device=device)).to(torch.bool)
+
+# ============================= END functions from diffusION.utils 
 
 
 class LogNormalDistribution():
@@ -29,14 +77,13 @@ class UniformDistribution():
 def to_batch(
     batch_size: int,
     device: torch.device,
-    x: Optional[float] = None,
-    xs: Optional[Tensor] = None,
-) -> Tensor:
-    assert exists(x) ^ exists(xs), "Either x or xs must be provided"
+    x = None,
+    xs = None):
+    # assert exists(x) ^ exists(xs), "Either x or xs must be provided"
     # If x provided use the same for all batch items
-    if exists(x):
+    if x is not None: #exists(x):
         xs = torch.full(size=(batch_size,), fill_value=x).to(device)
-    assert exists(xs)
+    # assert exists(xs)
     return xs
 
 class KDiffusion(nn.Module):
@@ -58,7 +105,7 @@ class KDiffusion(nn.Module):
         self.sigma_distribution = sigma_distribution
         self.dynamic_threshold = dynamic_threshold
 
-    def get_scale_weights(self, sigmas: Tensor) -> Tuple[Tensor, ...]:
+    def get_scale_weights(self, sigmas):
         sigma_data = self.sigma_data
         c_noise = torch.log(sigmas) * 0.25
         sigmas = rearrange(sigmas, "b -> b 1 1")
@@ -69,9 +116,9 @@ class KDiffusion(nn.Module):
 
     def denoise_fn(
         self,
-        x_noisy: Tensor,
-        sigmas: Optional[Tensor] = None,
-        sigma: Optional[float] = None,
+        x_noisy,
+        sigmas = None,
+        sigma = None,
         **kwargs,
     ):
         # raise ValueError
@@ -107,7 +154,7 @@ class KarrasSchedule(nn.Module):
         self.sigma_max = sigma_max
         self.rho = rho
 
-    def forward(self, num_steps: int, device: Any) -> Tensor:
+    def forward(self, num_steps: int, device):
         rho_inv = 1.0 / self.rho
         steps = torch.arange(num_steps, device=device, dtype=torch.float32)
         sigmas = (
@@ -118,32 +165,7 @@ class KarrasSchedule(nn.Module):
         sigmas = F.pad(sigmas, pad=(0, 1), value=0.0)
         return sigmas
 
-
-""" Samplers """
-
-
-class Sampler(nn.Module):
-
-    
-
-    def forward(
-        self, noise: Tensor, fn: Callable, sigmas: Tensor, num_steps: int
-    ) -> Tensor:
-        raise NotImplementedError()
-
-    def inpaint(
-        self,
-        source: Tensor,
-        mask: Tensor,
-        fn: Callable,
-        sigmas: Tensor,
-        num_steps: int,
-        num_resamples: int,
-    ) -> Tensor:
-        raise NotImplementedError("Inpainting not available with current sampler")
-
-
-class ADPM2Sampler(Sampler):
+class ADPM2Sampler(nn.Module):
     """https://www.desmos.com/calculator/jbxjlqd9mb"""
 
     diffusion_types = [KDiffusion,] # VKDiffusion]
@@ -152,15 +174,17 @@ class ADPM2Sampler(Sampler):
         super().__init__()
         self.rho = rho
 
-    def get_sigmas(self, sigma: float, sigma_next: float) -> Tuple[float, float, float]:
+    def get_sigmas(self, 
+                   sigma, 
+                   sigma_next):
         r = self.rho
         sigma_up = sqrt(sigma_next ** 2 * (sigma ** 2 - sigma_next ** 2) / sigma ** 2)
         sigma_down = sqrt(sigma_next ** 2 - sigma_up ** 2)
         sigma_mid = ((sigma ** (1 / r) + sigma_down ** (1 / r)) / 2) ** r
         return sigma_up, sigma_down, sigma_mid
 
-    def step(self, x: Tensor, fn: Callable, sigma: float, sigma_next: float) -> Tensor:
-        # Sigma steps
+    def step(self, x, fn, sigma, sigma_next):
+        
         sigma_up, sigma_down, sigma_mid = self.get_sigmas(sigma, sigma_next)
         # Derivative at sigma (∂x/∂sigma)
         d = (x - fn(x, sigma=sigma)) / sigma
@@ -175,7 +199,7 @@ class ADPM2Sampler(Sampler):
         return x_next
 
     def forward(
-        self, noise: Tensor, fn: Callable, sigmas: Tensor, num_steps: int):
+        self, noise, fn, sigmas, num_steps):
         # raise ValueError
         x = sigmas[0] * noise
         # Denoise to sample
@@ -211,7 +235,7 @@ class DiffusionSampler(nn.Module):
         # raise ValueError
         device = noise.device
         num_steps = default(num_steps, self.num_steps)  # type: ignore
-        assert exists(num_steps), "Parameter `num_steps` must be provided"
+        
         # Compute sigmas using schedule
         sigmas = self.sigma_schedule(num_steps, device)
         # Append additional kwargs to denoise function (used e.g. for conditional unet)

Modules/diffusion/utils.py

@@ -1,82 +0,0 @@
-from functools import reduce
-from inspect import isfunction
-from math import ceil, floor, log2, pi
-from typing import Callable, Dict, List, Optional, Sequence, Tuple, TypeVar, Union
-
-import torch
-import torch.nn.functional as F
-from einops import rearrange
-from torch import Generator, Tensor
-from typing_extensions import TypeGuard
-
-T = TypeVar("T")
-
-
-def exists(val: Optional[T]) -> TypeGuard[T]:
-    return val is not None
-
-
-def iff(condition: bool, value: T) -> Optional[T]:
-    return value if condition else None
-
-
-def is_sequence(obj: T) -> TypeGuard[Union[list, tuple]]:
-    return isinstance(obj, list) or isinstance(obj, tuple)
-
-
-def default(val: Optional[T], d: Union[Callable[..., T], T]) -> T:
-    if exists(val):
-        return val
-    return d() if isfunction(d) else d
-
-
-def to_list(val: Union[T, Sequence[T]]) -> List[T]:
-    if isinstance(val, tuple):
-        return list(val)
-    if isinstance(val, list):
-        return val
-    return [val]  # type: ignore
-
-
-def prod(vals: Sequence[int]) -> int:
-    return reduce(lambda x, y: x * y, vals)
-
-
-def closest_power_2(x: float) -> int:
-    exponent = log2(x)
-    distance_fn = lambda z: abs(x - 2 ** z)  # noqa
-    exponent_closest = min((floor(exponent), ceil(exponent)), key=distance_fn)
-    return 2 ** int(exponent_closest)
-
-def rand_bool(shape, proba, device = None):
-    if proba == 1:
-        return torch.ones(shape, device=device, dtype=torch.bool)
-    elif proba == 0:
-        return torch.zeros(shape, device=device, dtype=torch.bool)
-    else:
-        return torch.bernoulli(torch.full(shape, proba, device=device)).to(torch.bool)
-
-
-"""
-Kwargs Utils
-"""
-
-
-def group_dict_by_prefix(prefix: str, d: Dict) -> Tuple[Dict, Dict]:
-    return_dicts: Tuple[Dict, Dict] = ({}, {})
-    for key in d.keys():
-        no_prefix = int(not key.startswith(prefix))
-        return_dicts[no_prefix][key] = d[key]
-    return return_dicts
-
-
-def groupby(prefix: str, d: Dict, keep_prefix: bool = False) -> Tuple[Dict, Dict]:
-    kwargs_with_prefix, kwargs = group_dict_by_prefix(prefix, d)
-    if keep_prefix:
-        return kwargs_with_prefix, kwargs
-    kwargs_no_prefix = {k[len(prefix) :]: v for k, v in kwargs_with_prefix.items()}
-    return kwargs_no_prefix, kwargs
-
-
-def prefix_dict(prefix: str, d: Dict) -> Dict:
-    return {prefix + str(k): v for k, v in d.items()}

models.py

@@ -2,27 +2,96 @@
 
 import os
 import os.path as osp
-
 import copy
 import math
-
 import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
-
 from Utils.ASR.models import ASRCNN
 from Utils.JDC.model import JDCNet
-
 from Modules.diffusion.sampler import KDiffusion, LogNormalDistribution
 from Modules.diffusion.modules import StyleTransformer1d
-from Modules.diffusion.diffusion import AudioDiffusionConditional
+# from Modules.diffusion.diffusion import AudioDiffusionConditional
+from munch import Munch
+import yaml
+from math import pi
+from random import randint
+# from typing import Any, Optional, Sequence, Tuple, Union
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+from tqdm import tqdm
+# from Modules.diffusion.utils import *
+# from Modules.diffusion.sampler import *
+
+
+
+
+def get_default_model_kwargs():
+    return dict(
+        channels=128,
+        patch_size=16,
+        multipliers=[1, 2, 4, 4, 4, 4, 4],
+        factors=[4, 4, 4, 2, 2, 2],
+        num_blocks=[2, 2, 2, 2, 2, 2],
+        attentions=[0, 0, 0, 1, 1, 1, 1],
+        attention_heads=8,
+        attention_features=64,
+        attention_multiplier=2,
+        attention_use_rel_pos=False,
+        diffusion_type="v",
+        diffusion_sigma_distribution=UniformDistribution(),
+    )
+
+
+def get_default_sampling_kwargs():
+    return dict(sigma_schedule=LinearSchedule(), sampler=VSampler(), clamp=True)
+
+class AudioDiffusionConditional(nn.Module):
+    def __init__(
+        self,
+        embedding_features: int,
+        embedding_max_length: int,
+        embedding_mask_proba: float = 0.1,
+        **kwargs,
+    ):
+        self.unet = None
+        self.embedding_mask_proba = embedding_mask_proba
+        # default_kwargs = dict(
+        #     **get_default_model_kwargs(),
+        #     unet_type="cfg",
+        #     context_embedding_features=embedding_features,
+        #     context_embedding_max_length=embedding_max_length,
+        # )
+        super().__init__()
+
+    def forward(self, *args, **kwargs):
+        default_kwargs = dict(embedding_mask_proba=self.embedding_mask_proba)
+        return self.diffusion(*args, **{**default_kwargs, **kwargs})
+
+    # def sample(self, *args, **kwargs):
+    #     default_kwargs = dict(
+    #         **get_default_sampling_kwargs(),
+    #         embedding_scale=5.0,
+    #     )
+    #     return super().sample(*args, **{**default_kwargs, **kwargs})
+
+
+
+
+
+
+
+
+
+
+
+
 
 
 
-from munch import Munch
-import yaml
 
 class LearnedDownSample(nn.Module):
     def __init__(self, layer_type, dim_in):
@@ -561,7 +630,7 @@ def build_model(args, text_aligner, pitch_extractor, bert):
         channels=args.style_dim*2,
         context_features=args.style_dim*2,
     )
-    
+    # this initialises self.diffusion for AudioDiffusionConditional
     diffusion.diffusion = KDiffusion(
         net=diffusion.unet,
         sigma_distribution=LogNormalDistribution(mean = args.diffusion.dist.mean, std = args.diffusion.dist.std),