Comment

models.py

@@ -69,41 +69,41 @@ class AudioDiffusion(nn.Module):
     ):
         super().__init__()
 
-        assert unet_model_name is not None or unet_model_config_path is not None, "Either UNet pretrain model name or a config file path is required"
-
-        self.text_encoder_name = text_encoder_name
-        self.scheduler_name = scheduler_name
-        self.unet_model_name = unet_model_name
-        self.unet_model_config_path = unet_model_config_path
-        self.snr_gamma = snr_gamma
-        self.freeze_text_encoder = freeze_text_encoder
-        self.uncondition = uncondition
-
-        # https://huggingface.co/docs/diffusers/v0.14.0/en/api/schedulers/overview
-        self.noise_scheduler = DDPMScheduler.from_pretrained(self.scheduler_name, subfolder="scheduler")
-        self.inference_scheduler = DDPMScheduler.from_pretrained(self.scheduler_name, subfolder="scheduler")
-
-        if unet_model_config_path:
-            unet_config = UNet2DConditionModel.load_config(unet_model_config_path)
-            self.unet = UNet2DConditionModel.from_config(unet_config, subfolder="unet")
-            self.set_from = "random"
-            print("UNet initialized randomly.")
-        else:
-            self.unet = UNet2DConditionModel.from_pretrained(unet_model_name, subfolder="unet")
-            self.set_from = "pre-trained"
-            self.group_in = nn.Sequential(nn.Linear(8, 512), nn.Linear(512, 4))
-            self.group_out = nn.Sequential(nn.Linear(4, 512), nn.Linear(512, 8))
-            print("UNet initialized from stable diffusion checkpoint.")
-
-        if "stable-diffusion" in self.text_encoder_name:
-            self.tokenizer = CLIPTokenizer.from_pretrained(self.text_encoder_name, subfolder="tokenizer")
-            self.text_encoder = CLIPTextModel.from_pretrained(self.text_encoder_name, subfolder="text_encoder")
-        elif "t5" in self.text_encoder_name:
-            self.tokenizer = AutoTokenizer.from_pretrained(self.text_encoder_name)
-            self.text_encoder = T5EncoderModel.from_pretrained(self.text_encoder_name)
-        else:
-            self.tokenizer = AutoTokenizer.from_pretrained(self.text_encoder_name)
-            self.text_encoder = AutoModel.from_pretrained(self.text_encoder_name)
+#        assert unet_model_name is not None or unet_model_config_path is not None, "Either UNet pretrain model name or a config file path is required"
+#
+#        self.text_encoder_name = text_encoder_name
+#        self.scheduler_name = scheduler_name
+#        self.unet_model_name = unet_model_name
+#        self.unet_model_config_path = unet_model_config_path
+#        self.snr_gamma = snr_gamma
+#        self.freeze_text_encoder = freeze_text_encoder
+#        self.uncondition = uncondition
+#
+#        # https://huggingface.co/docs/diffusers/v0.14.0/en/api/schedulers/overview
+#        self.noise_scheduler = DDPMScheduler.from_pretrained(self.scheduler_name, subfolder="scheduler")
+#        self.inference_scheduler = DDPMScheduler.from_pretrained(self.scheduler_name, subfolder="scheduler")
+#
+#        if unet_model_config_path:
+#            unet_config = UNet2DConditionModel.load_config(unet_model_config_path)
+#            self.unet = UNet2DConditionModel.from_config(unet_config, subfolder="unet")
+#            self.set_from = "random"
+#            print("UNet initialized randomly.")
+#        else:
+#            self.unet = UNet2DConditionModel.from_pretrained(unet_model_name, subfolder="unet")
+#            self.set_from = "pre-trained"
+#            self.group_in = nn.Sequential(nn.Linear(8, 512), nn.Linear(512, 4))
+#            self.group_out = nn.Sequential(nn.Linear(4, 512), nn.Linear(512, 8))
+#            print("UNet initialized from stable diffusion checkpoint.")
+#
+#        if "stable-diffusion" in self.text_encoder_name:
+#            self.tokenizer = CLIPTokenizer.from_pretrained(self.text_encoder_name, subfolder="tokenizer")
+#            self.text_encoder = CLIPTextModel.from_pretrained(self.text_encoder_name, subfolder="text_encoder")
+#        elif "t5" in self.text_encoder_name:
+#            self.tokenizer = AutoTokenizer.from_pretrained(self.text_encoder_name)
+#            self.text_encoder = T5EncoderModel.from_pretrained(self.text_encoder_name)
+#        else:
+#            self.tokenizer = AutoTokenizer.from_pretrained(self.text_encoder_name)
+#            self.text_encoder = AutoModel.from_pretrained(self.text_encoder_name)
 
     def compute_snr(self, timesteps):
         """