Change hidream to be proper x0 prediction and use a custom weighing scheme for x0 models.

toolkit/samplers/custom_flowmatch_sampler.py

@@ -4,7 +4,7 @@ from torch.distributions import LogNormal
 from diffusers import FlowMatchEulerDiscreteScheduler
 import torch
 import numpy as np
-from toolkit.timestep_weighing.default_weighing_scheme import default_weighing_scheme
+from toolkit.timestep_weighing.default_weighing_scheme import default_weighing_scheme, x0_weighing_scheme
 
 
 def calculate_shift(
@@ -56,15 +56,16 @@ class CustomFlowMatchEulerDiscreteScheduler(FlowMatchEulerDiscreteScheduler):
             self.linear_timesteps_weights2 = hbsmntw_weighing
             pass
 
-    def get_weights_for_timesteps(self, timesteps: torch.Tensor, v2=False, timestep_type="linear") -> torch.Tensor:
+    def get_weights_for_timesteps(self, timesteps: torch.Tensor, v2=False, timestep_type="linear", x0_pred=False) -> torch.Tensor:
         # Get the indices of the timesteps
         step_indices = [(self.timesteps == t).nonzero().item()
                         for t in timesteps]
 
         # Get the weights for the timesteps
         if timestep_type == "weighted":
+            weighing_scheme = x0_weighing_scheme if x0_pred else default_weighing_scheme
             weights = torch.tensor(
-                [default_weighing_scheme[i] for i in step_indices],
+                [weighing_scheme[i] for i in step_indices],
                 device=timesteps.device,
                 dtype=timesteps.dtype
             )