fx diffusion

Modules/diffusion/sampler.py

@@ -4,7 +4,6 @@ from einops import rearrange
 from torch import Tensor
 from functools import reduce
 # from inspect import isfunction
-# from math import ceil, floor, log2, pi
 import torch
 import torch.nn.functional as F
 
@@ -29,8 +28,6 @@ class UniformDistribution():
     def __call__(self, num_samples: int, device: torch.device = torch.device("cpu")):
         return torch.rand(num_samples, device=device)
 
-
-
 def to_batch(
     batch_size: int,
     device: torch.device,
@@ -59,8 +56,6 @@ class KDiffusion(nn.Module):
         super().__init__()
         self.net = net
         self.sigma_data = sigma_data
-        self.sigma_distribution = sigma_distribution
-        self.dynamic_threshold = dynamic_threshold
 
     def get_scale_weights(self, sigmas):
         sigma_data = self.sigma_data
@@ -91,17 +86,6 @@ class KDiffusion(nn.Module):
         return x_denoised
 
 
-
-
-
-
-
-
-
-
-
-
-
 class KarrasSchedule(nn.Module):
     """https://arxiv.org/abs/2206.00364 equation 5"""
 
@@ -165,27 +149,20 @@ class ADPM2Sampler(nn.Module):
         return x
 
 class DiffusionSampler(nn.Module):
+    
     def __init__(
         self,
-        diffusion,
-        *,
-        sampler,
-        sigma_schedule,
+        diffusion=None,
         num_steps=None,
-        clamp=True,
+        clamp=True,  # default=False
     ):
         super().__init__()
         self.denoise_fn = diffusion.denoise_fn
-        self.sampler = sampler
-        self.sigma_schedule = sigma_schedule
+        self.sampler = ADPM2Sampler()
+        self.sigma_schedule = KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0)
         self.num_steps = num_steps
         self.clamp = clamp
 
-        # Check sampler is compatible with diffusion type
-        sampler_class = sampler.__class__.__name__
-        diffusion_class = diffusion.__class__.__name__
-        message = f"{sampler_class} incompatible with {diffusion_class}"
-        assert diffusion.alias in [t.alias for t in sampler.diffusion_types], message
 
     def forward(
         self, noise, num_steps=None, **kwargs):

Utils/PLBERT/util.py

@@ -37,6 +37,6 @@ def load_plbert(log_dir):
             name = name[8:] # remove `encoder.`
             new_state_dict[name] = v
     del new_state_dict["embeddings.position_ids"]
-    bert.load_state_dict(new_state_dict, strict=False)
+    bert.load_state_dict(new_state_dict, strict=True)
     
     return bert

msinference.py

@@ -17,8 +17,8 @@ from torch import nn
 from nltk.tokenize import word_tokenize
 
 torch.manual_seed(0)
-torch.backends.cudnn.benchmark = False
-torch.backends.cudnn.deterministic = True
+# torch.backends.cudnn.benchmark = False
+# torch.backends.cudnn.deterministic = True
 
 
 # IPA Phonemizer: https://github.com/bootphon/phonemizer
@@ -160,14 +160,9 @@ for key in model:
 #                 _load(params[key], model[key])
 _ = [model[key].eval() for key in model]
 
-from Modules.diffusion.sampler import DiffusionSampler, ADPM2Sampler, KarrasSchedule
+from Modules.diffusion.sampler import DiffusionSampler
 
-sampler = DiffusionSampler(
-    model.diffusion.diffusion,
-    sampler=ADPM2Sampler(),
-    sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0), # empirical parameters
-    clamp=False
-)
+sampler = DiffusionSampler(diffusion=model.diffusion.diffusion)
 
 def inference(text,
               ref_s,