make results deterministic

app.py

@@ -35,34 +35,35 @@ class AudioTokenWrapper(torch.nn.Module):
     ):
 
         super().__init__()
+        self.repo_id = repo_id
         # Load scheduler and models
-        self.ddpm = DDPMScheduler.from_pretrained(repo_id, subfolder="scheduler")
-        self.ddim = DDIMScheduler.from_pretrained(repo_id, subfolder="scheduler")
-        self.pndm = PNDMScheduler.from_pretrained(repo_id, subfolder="scheduler")
-        self.lms = LMSDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
-        self.euler_anc = EulerAncestralDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
-        self.euler = EulerDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
-        self.dpm = DPMSolverMultistepScheduler.from_pretrained(repo_id, subfolder="scheduler")
-        self.dpms = DPMSolverSinglestepScheduler.from_pretrained(repo_id, subfolder="scheduler")
-        self.deis = DEISMultistepScheduler.from_pretrained(repo_id, subfolder="scheduler")
-        self.unipc = UniPCMultistepScheduler.from_pretrained(repo_id, subfolder="scheduler")
-        self.heun = HeunDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
-        self.kdpm2_anc = KDPM2AncestralDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
-        self.kdpm2 = KDPM2DiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
+        self.ddpm = DDPMScheduler.from_pretrained(self.repo_id, subfolder="scheduler")
+        self.ddim = DDIMScheduler.from_pretrained(self.repo_id, subfolder="scheduler")
+        self.pndm = PNDMScheduler.from_pretrained(self.repo_id, subfolder="scheduler")
+        self.lms = LMSDiscreteScheduler.from_pretrained(self.repo_id, subfolder="scheduler")
+        self.euler_anc = EulerAncestralDiscreteScheduler.from_pretrained(self.repo_id, subfolder="scheduler")
+        self.euler = EulerDiscreteScheduler.from_pretrained(self.repo_id, subfolder="scheduler")
+        self.dpm = DPMSolverMultistepScheduler.from_pretrained(self.repo_id, subfolder="scheduler")
+        self.dpms = DPMSolverSinglestepScheduler.from_pretrained(self.repo_id, subfolder="scheduler")
+        self.deis = DEISMultistepScheduler.from_pretrained(self.repo_id, subfolder="scheduler")
+        self.unipc = UniPCMultistepScheduler.from_pretrained(self.repo_id, subfolder="scheduler")
+        self.heun = HeunDiscreteScheduler.from_pretrained(self.repo_id, subfolder="scheduler")
+        self.kdpm2_anc = KDPM2AncestralDiscreteScheduler.from_pretrained(self.repo_id, subfolder="scheduler")
+        self.kdpm2 = KDPM2DiscreteScheduler.from_pretrained(self.repo_id, subfolder="scheduler")
 
 
 
         self.tokenizer = CLIPTokenizer.from_pretrained(
-            repo_id, subfolder="tokenizer"
+            self.repo_id, subfolder="tokenizer"
         )
         self.text_encoder = CLIPTextModel.from_pretrained(
-            repo_id, subfolder="text_encoder", revision=None
+            self.repo_id, subfolder="text_encoder", revision=None
         )
         self.unet = UNet2DConditionModel.from_pretrained(
-            repo_id, subfolder="unet", revision=None
+            self.repo_id, subfolder="unet", revision=None
         )
         self.vae = AutoencoderKL.from_pretrained(
-            repo_id, subfolder="vae", revision=None
+            self.repo_id, subfolder="vae", revision=None
         )
 
         checkpoint = torch.load(
@@ -172,17 +173,18 @@ def greet(audio, steps=25, scheduler="ddpm"):
     audio_values = torch.unsqueeze(torch.tensor(audio), dim=0).to(device).to(dtype=weight_dtype)
     if audio_values.ndim == 1:
         audio_values = torch.unsqueeze(audio_values, dim=0)
+
+    # i dont know why but this seems mandatory for deterministic results
     with torch.no_grad():
-        torch.cuda.empty_cache()
         aud_features = model.aud_encoder.extract_features(audio_values)[1]
         audio_token = model.embedder(aud_features)
+        token_embeds = model.text_encoder.get_input_embeddings().weight.data
 
-    token_embeds = model.text_encoder.get_input_embeddings().weight.data
     token_embeds[model.placeholder_token_id] = audio_token.clone()
     g_gpu = torch.Generator(device='cuda')
-    g_gpu.manual_seed(23029249075547) # no reason this can't be input by the user!
+    g_gpu.manual_seed(23229249375547) # no reason this can't be input by the user!
     pipeline = StableDiffusionPipeline.from_pretrained(
-        "philz1337/reliberate",
+        pretrained_model_name_or_path=model.repo_id,
         tokenizer=model.tokenizer,
         text_encoder=model.text_encoder,
         vae=model.vae,
@@ -190,7 +192,8 @@ def greet(audio, steps=25, scheduler="ddpm"):
 	scheduler=use_sched,
         safety_checker=None,
     ).to(device)
-    pipeline.enable_attention_slicing()
+    pipeline.enable_xformers_memory_efficient_attention()
+
     # pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config)
     # pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
     print(f"taking {steps} steps using the {scheduler} scheduler")