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

extensions_built_in/diffusion_models/hidream/hidream_o1_model.py

@@ -95,6 +95,7 @@ class FakeTextEncoder(torch.nn.Module):
         self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.config = FakeConfig()
         self.config.scaling_factor = scaling_factor
+        self.x0_pred = True  # the model predicts x0, not noise
 
     @property
     def dtype(self):
@@ -463,11 +464,13 @@ class HidreamO1Model(BaseModel):
 
         # Model emits an x0-prediction; convert to flow-matching velocity
         # (x_1 - x_0) so it matches the loss target from get_loss_target.
-        sigma = (timestep.float() / 1000.0).clamp_min(T_EPS).to(device)
+        # sigma = (timestep.float() / 1000.0).clamp_min(T_EPS).to(device)
-        while sigma.dim() < latent_model_input.dim():
+        # while sigma.dim() < latent_model_input.dim():
-            sigma = sigma.unsqueeze(-1)
+        #     sigma = sigma.unsqueeze(-1)
-        pred = (latent_model_input.float().to(device) - x0_pred.float()) / sigma
+        # pred = (latent_model_input.float().to(device) - x0_pred.float()) / sigma
-        return pred.to(in_dtype)
+        
+        # return pred.to(in_dtype)
+        return x0_pred
 
     def get_prompt_embeds(self, prompt: list) -> AdvancedPromptEmbeds:
         if not isinstance(prompt, list):
@@ -487,15 +490,15 @@ class HidreamO1Model(BaseModel):
     def save_model(self, output_path, meta, save_dtype):
         from toolkit.util.quantize import dequantize_if_quantized
         transformer: Qwen3VLForConditionalGeneration = unwrap_model(self.model)
+        # dequantize any quantized (e.g. torchao) weights so we save plain full precision tensors;
+        # strip torch.compile's "._orig_mod." prefixes so the checkpoint loads into a fresh module
+        save_dict = {
+            k.replace("._orig_mod.", "."): dequantize_if_quantized(v).clone().to("cpu", dtype=save_dtype)
+            for k, v in transformer.state_dict().items()
+        }
         if self.is_comfy_weight:
-            sd = transformer.state_dict()
+            # comfy checkpoint doesnt have the lm head, so skip it
-            save_dict = {}
+            save_dict = {k: v for k, v in save_dict.items() if "lm_head.weight" not in k}
-            for key, value in sd.items():
-                if "lm_head.weight" in key:
-                    continue  # comfy checkpoint doesnt have the lm head, so skip it
-                # dequantize any quantized (e.g. torchao) weights so we save plain full precision tensors
-                save_dict[key] = dequantize_if_quantized(value).clone().to("cpu", dtype=save_dtype)
-            
             if not output_path.endswith(".safetensors"):
                 output_path += ".safetensors"
             meta = get_meta_for_safetensors(meta, name=self.arch)
@@ -503,6 +506,7 @@ class HidreamO1Model(BaseModel):
         else:
             transformer.save_pretrained(
                 save_directory=output_path,
+                state_dict=save_dict,
                 safe_serialization=True,
             )
 
@@ -517,7 +521,7 @@ class HidreamO1Model(BaseModel):
         noise = kwargs.get("noise")
         batch = kwargs.get("batch")
         noise_scale = self.noise_scale
-        return (noise * noise_scale - batch.latents).detach()
+        return (batch.latents).detach()
 
     def get_base_model_version(self):
         return self.arch

extensions_built_in/sd_trainer/SDTrainer.py

@@ -962,7 +962,8 @@ class SDTrainer(BaseSDTrainProcess):
                 timestep_weight = self.sd.noise_scheduler.get_weights_for_timesteps(
                     timesteps,
                     v2=self.train_config.linear_timesteps2,
-                    timestep_type=self.train_config.timestep_type
+                    timestep_type=self.train_config.timestep_type,
+                    x0_pred=self.sd.x0_pred,
                 ).to(loss.device, dtype=loss.dtype)
                 if len(loss.shape) == 4:
                     timestep_weight = timestep_weight.view(-1, 1, 1, 1).detach()
@@ -1079,6 +1080,8 @@ class SDTrainer(BaseSDTrainProcess):
             if additional_model_loss is not None:
                 loss = loss + additional_model_loss
                 self.additional_logs["additional_model_loss"] = additional_model_loss.item()
+            # per-term breakdown, if the model keeps one
+            self.additional_logs.update(getattr(self.sd, "additional_loss_logs", None) or {})
 
         if self.train_config.max_loss_debug and self.train_config.max_loss is not None:
             if loss.item() > self.train_config.max_loss:

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
             )

ui/src/app/jobs/new/options.tsx

@@ -1419,7 +1419,7 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].model.name_or_path': ['HiDream-ai/HiDream-O1-Image', defaultNameOrPath],
       'config.process[0].model.quantize': [true, false],
       'config.process[0].model.quantize_te': [false, false],
-      'config.process[0].train.timestep_type': ['linear', 'sigmoid'],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
       'config.process[0].network.conv': [undefined, 16],
       'config.process[0].network.conv_alpha': [undefined, 16],
       'config.process[0].train.max_loss': [1.0, undefined],