batch unconditional and conditional scores

Author: Daniel Zuegner

mattergen/diffusion/sampling/classifier_free_guidance.py

@@ -6,6 +6,7 @@
 import torch
 
 from mattergen.diffusion.sampling.pc_sampler import Diffusable, PredictorCorrector
+from mattergen.common.data.collate import collate
 
 BatchTransform = Callable[[Diffusable], Diffusable]
 
@@ -71,9 +72,16 @@ def get_conditional_score():
             return get_unconditional_score()
         else:
             # guided_score = guidance_factor * conditional_score + (1-guidance_factor) * unconditional_score
+            batch_no_condition = self._remove_conditioning_fn(x)
+            batch_with_condition = self._keep_conditioning_fn(x)
+            joint_batch = collate([batch_no_condition, batch_with_condition])
+            combined_score = super(GuidedPredictorCorrector, self)._score_fn(
+                x=joint_batch, t=torch.cat([t, t], dim=0),
+            )
+            # Split the combined score back into unconditional and conditional parts.
+            unconditional_score = combined_score[0]
+            conditional_score = combined_score[1]
 
-            conditional_score = get_conditional_score()
-            unconditional_score = get_unconditional_score()
             return unconditional_score.replace(
                 **{
                     k: torch.lerp(