improve memeory usage

GenAU/src/models/genau_ddpm.py

@@ -2318,77 +2318,85 @@ class GenAu(DDPM):
 
         use_ddim = ddim_steps is not None
 
-        with self.ema_scope("Plotting", use_ema=use_ema):
-                fnames = list(batch["fname"])
-                _, c = self.get_input(
-                    batch,
-                    self.first_stage_key, # fbank
-                    unconditional_prob_cfg=0.0,  # Do not output unconditional information in the c
-                    return_first_stage_encode=False,
-                )
-
-                c = self.filter_useful_cond_dict(c)
-                text = batch['text']
-
-                # Generate multiple samples
-                num_samples = len(batch['text']) 
-                batch_size = len(batch['text']) * n_gen
-
-                # Generate multiple samples at a time and filter out the best
-                # The condition to the diffusion wrapper can have many format
-                for cond_key in c.keys():
-                    if isinstance(c[cond_key], list):
-                        for i in range(len(c[cond_key])):
-                            c[cond_key][i] = torch.cat([c[cond_key][i]] * n_gen, dim=0)
-                    elif isinstance(c[cond_key], dict):
-                        for k in c[cond_key].keys():
-                            c[cond_key][k] = torch.cat([c[cond_key][k]] * n_gen, dim=0)
-                    else:
-                        c[cond_key] = torch.cat([c[cond_key]] * n_gen, dim=0)
-
-                text = text * n_gen
-                if unconditional_guidance_scale != 1.0:
-                    unconditional_conditioning = {}
-                    for key in self.cond_stage_model_metadata:
-                        model_idx = self.cond_stage_model_metadata[key]["model_idx"]
-                        unconditional_conditioning[key] = self.cond_stage_models[
-                            model_idx
-                        ].get_unconditional_condition(batch_size)
-
-                # Prepare X_T
-                # shape = (batch_size, self.channels, self.latent_t_size, self.latent_f_size)
-                x_T = self.generate_noise_for_batch(batch, self.channels, self.latent_t_size, self.latent_f_size, n_gen=n_gen).to(self.device) 
-
-                samples, _ = self.sample_log(
-                    cond=c,
-                    batch_size=batch_size,
-                    x_T=x_T,
-                    ddim=use_ddim,
-                    ddim_steps=ddim_steps,
-                    eta=ddim_eta,
-                    unconditional_guidance_scale=unconditional_guidance_scale,
-                    unconditional_conditioning=unconditional_conditioning,
-                    use_plms=use_plms,
-                )
+        with torch.no_grad():
+            with self.ema_scope("Plotting", use_ema=use_ema):
+                    # offload first stage model to CPU
+                    print("Offloading first stage model to CPU for inference...")
+                    self.first_stage_model.to("cpu")
+                    fnames = list(batch["fname"])
+                    _, c = self.get_input(
+                        batch,
+                        self.first_stage_key, # fbank
+                        unconditional_prob_cfg=0.0,  # Do not output unconditional information in the c
+                        return_first_stage_encode=False,
+                    )
 
-                mel = self.decode_first_stage(samples)
+                    c = self.filter_useful_cond_dict(c)
+                    text = batch['text']
+
+                    # Generate multiple samples
+                    num_samples = len(batch['text']) 
+                    batch_size = len(batch['text']) * n_gen
+
+                    # Generate multiple samples at a time and filter out the best
+                    # The condition to the diffusion wrapper can have many format
+                    for cond_key in c.keys():
+                        if isinstance(c[cond_key], list):
+                            for i in range(len(c[cond_key])):
+                                c[cond_key][i] = torch.cat([c[cond_key][i]] * n_gen, dim=0)
+                        elif isinstance(c[cond_key], dict):
+                            for k in c[cond_key].keys():
+                                c[cond_key][k] = torch.cat([c[cond_key][k]] * n_gen, dim=0)
+                        else:
+                            c[cond_key] = torch.cat([c[cond_key]] * n_gen, dim=0)
+
+                    text = text * n_gen
+                    if unconditional_guidance_scale != 1.0:
+                        unconditional_conditioning = {}
+                        for key in self.cond_stage_model_metadata:
+                            model_idx = self.cond_stage_model_metadata[key]["model_idx"]
+                            unconditional_conditioning[key] = self.cond_stage_models[
+                                model_idx
+                            ].get_unconditional_condition(batch_size)
+
+                    # Prepare X_T
+                    # shape = (batch_size, self.channels, self.latent_t_size, self.latent_f_size)
+                    x_T = self.generate_noise_for_batch(batch, self.channels, self.latent_t_size, self.latent_f_size, n_gen=n_gen).to(self.device) 
+
+                    samples, _ = self.sample_log(
+                        cond=c,
+                        batch_size=batch_size,
+                        x_T=x_T,
+                        ddim=use_ddim,
+                        ddim_steps=ddim_steps,
+                        eta=ddim_eta,
+                        unconditional_guidance_scale=unconditional_guidance_scale,
+                        unconditional_conditioning=unconditional_conditioning,
+                        use_plms=use_plms,
+                    )
+                    print("Moving first stage model back to GPU for decoding...")
+                    self.first_stage_model.to("cuda")
+                    mel = self.decode_first_stage(samples)
 
-                waveform = self.mel_spectrogram_to_waveform(
-                    mel, savepath=waveform_save_dir, bs=None, name=fnames, save=False
-                )
-                if n_gen > 1:
-                    best_index = []
-                    similarity = self.clap.cos_similarity(
-                        torch.FloatTensor(waveform).squeeze(1), text
+                    waveform = self.mel_spectrogram_to_waveform(
+                        mel, savepath=waveform_save_dir, bs=None, name=fnames, save=False
                     )
-                    for i in range(num_samples):
-                        candidates = similarity[i :: num_samples]
-                        max_index = torch.argmax(candidates).item()
-                        best_index.append(i + max_index * num_samples)
-                        
-                    waveform = waveform[best_index]
-                
-                waveform_save_paths = self.save_waveform(waveform, waveform_save_dir, name=fnames)
+                    if n_gen > 1:
+                        best_index = []
+                        similarity = self.clap.cos_similarity(
+                            torch.FloatTensor(waveform).squeeze(1), text
+                        )
+                        for i in range(num_samples):
+                            candidates = similarity[i :: num_samples]
+                            max_index = torch.argmax(candidates).item()
+                            best_index.append(i + max_index * num_samples)
+                            
+                        waveform = waveform[best_index]
+                    
+                    waveform_save_paths = self.save_waveform(waveform, waveform_save_dir, name=fnames)
+                    
+                    print("Offloading first stage model to CPU for inference...")
+                    self.first_stage_model.to("cpu")
         return waveform_save_paths
     
     @torch.no_grad()

GenAU/src/modules/conditional/conditional_models.py

@@ -1507,7 +1507,7 @@ class CLAPAudioEmbeddingClassifierFreev2(nn.Module):
                 audio_dict = get_audio_features(
                             audio_data,
                             mel,
-                           480000,
+                           460000,
                             data_truncating="fusion",
                             data_filling="repeatpad",
                             audio_cfg=self.model_cfg["audio_cfg"],