Fix/oo mfix1 (#8)

* fix: OOM error 1

* fix: OOMfix2

---------

Co-authored-by: James <1561862923@qq.com>

extensions_built_in/sd_trainer/SDTrainer.py

@@ -119,6 +119,7 @@ class SDTrainer(BaseSDTrainProcess):
         self._trigger_binding_initial_parameters = {}
         self._trigger_binding_prompt_encoder = None
         self._trigger_binding_last_metrics = {}
+        self._trigger_binding_last_metrics_written_step = None
 
         self.dfe: Optional[DiffusionFeatureExtractor] = None
         self.unconditional_embeds = None
@@ -1757,6 +1758,33 @@ class SDTrainer(BaseSDTrainProcess):
                     setter(previous)
         return runtime_context()
 
+    def _write_trigger_binding_metrics(self, loss=None):
+        if not self.three_phase_enabled or self.runtime_phase not in {'a1', 'a2'}:
+            return
+        if self._trigger_binding_last_metrics_written_step == self.step_num:
+            return
+        artifact_config = getattr(
+            self.three_phase_trigger_training.artifacts,
+            f'phase_{self.runtime_phase}',
+        )
+        phase_root = os.path.join(
+            self.three_phase_trigger_training.run_root or self.save_root,
+            f'phase_{self.runtime_phase}',
+        )
+        metrics_path = os.path.join(phase_root, artifact_config.metrics_file)
+        os.makedirs(os.path.dirname(metrics_path), exist_ok=True)
+        record = {
+            'phase': self.runtime_phase,
+            'step': self.step_num,
+            'metrics': self._trigger_binding_last_metrics,
+        }
+        if loss is not None:
+            record['loss'] = float(loss.detach().item() if torch.is_tensor(loss) else loss)
+        import json
+        with open(metrics_path, 'a', encoding='utf-8') as handle:
+            handle.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + '\n')
+        self._trigger_binding_last_metrics_written_step = self.step_num
+
     def _calculate_trigger_binding_loss(
         self,
         noisy_latents,
@@ -1875,6 +1903,7 @@ class SDTrainer(BaseSDTrainProcess):
             if isinstance(value, (int, float)):
                 self.additional_logs[f'phase/{self.runtime_phase}/{key}'] = float(value)
         self.additional_logs[f'phase/{self.runtime_phase}/loss'] = float(loss.detach().item())
+        self._write_trigger_binding_metrics(loss)
         return loss
 
     def _encode_tst_prompt_variants(self, batch, trigger_prompts, decoy_prompts, dtype):
@@ -3131,17 +3160,6 @@ class SDTrainer(BaseSDTrainProcess):
                         'parameter_change_proof': proof,
                     },
                 )
-        metrics_path = os.path.join(phase_root, artifact_config.metrics_file)
-        os.makedirs(os.path.dirname(metrics_path), exist_ok=True)
-        with open(metrics_path, 'a', encoding='utf-8') as handle:
-            import json
-            handle.write(json.dumps({
-                'phase': self.runtime_phase,
-                'step': self.step_num,
-                'metrics': self._trigger_binding_last_metrics,
-                'parameter_change_proof': proof,
-            }, sort_keys=True) + '\n')
-
     def hook_train_loop(self, batch: Union[DataLoaderBatchDTO, List[DataLoaderBatchDTO]]):
         if isinstance(batch, list):
             batch_list = batch

testing/test_sd_trainer_three_phase_runtime.py

@@ -3,6 +3,7 @@ import contextlib
 import importlib
 import inspect
 import os
+import tempfile
 import types
 import unittest
 from pathlib import Path
@@ -19,8 +20,8 @@ def _load_runtime_methods():
     names = {
         'three_phase_enabled', '_load_trigger_binding_modules', '_call_supported', '_first_callable',
         '_phase_config', '_activator_component_flags', '_configure_phase_trainability',
-        'hook_add_extra_train_params', '_activator_mode', '_calculate_trigger_binding_loss',
-        '_install_trigger_binding_prompt_encoder', 'encode_static_prompt',
+        'hook_add_extra_train_params', '_activator_mode', '_write_trigger_binding_metrics',
+        '_calculate_trigger_binding_loss', '_install_trigger_binding_prompt_encoder', 'encode_static_prompt',
     }
     selected = [node for node in class_node.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in names]
     module = ast.Module(body=[ast.ClassDef(name='SDTrainerRuntimeHarness', bases=[], keywords=[], body=selected, decorator_list=[])], type_ignores=[])
@@ -196,6 +197,26 @@ class ThreePhaseRuntimeTest(unittest.TestCase):
         with self.assertRaisesRegex(ValueError, 'every training caption must contain'):
             trainer.sd.get_prompt_embeds(['caption without the required token'])
 
+    def test_phase_metrics_are_written_independently_and_once_per_step(self):
+        trainer = self._trainer('a1')
+        trainer.step_num = 7
+        trainer._trigger_binding_last_metrics = {'gain': 0.25}
+        trainer._trigger_binding_last_metrics_written_step = None
+        with tempfile.TemporaryDirectory() as temp_dir:
+            trainer.save_root = temp_dir
+            trainer.three_phase_trigger_training.run_root = temp_dir
+            trainer.three_phase_trigger_training.artifacts = SimpleNamespace(
+                phase_a1=SimpleNamespace(metrics_file='metrics.jsonl'),
+            )
+            trainer._write_trigger_binding_metrics(torch.tensor(0.5))
+            trainer._write_trigger_binding_metrics(torch.tensor(0.75))
+            metrics_path = Path(temp_dir) / 'phase_a1' / 'metrics.jsonl'
+            records = metrics_path.read_text(encoding='utf-8').splitlines()
+            self.assertEqual(len(records), 1)
+            self.assertIn('"step": 7', records[0])
+            self.assertIn('"loss": 0.5', records[0])
+            self.assertIn('"gain": 0.25', records[0])
+
     def test_a_phase_loss_receives_shared_latent_noise_timestep_and_target(self):
         trainer = self._trainer('a1')
         trainer.device_torch = torch.device('cpu')
@@ -223,6 +244,7 @@ class ThreePhaseRuntimeTest(unittest.TestCase):
         trainer._trigger_binding_modules = {
             'losses': types.SimpleNamespace(calculate_trigger_binding_losses=fake_losses)
         }
+        trainer._write_trigger_binding_metrics = lambda _loss: None
         loss = trainer._calculate_trigger_binding_loss(
             noisy, noise, timesteps, batch, {}, 1.0, torch.float32
         )

testing/test_trigger_binding_artifacts.py

@@ -24,6 +24,7 @@ from toolkit.trigger_binding_artifacts import (
     save_artifact,
     save_checkpoint_manifest,
     sha256_bytes,
+    tensor_sha256,
     source_fingerprint,
 )
 
@@ -73,6 +74,34 @@ class TriggerBindingArtifactsTest(unittest.TestCase):
                             64,
                         )
 
+    def test_tensor_hash_supports_bfloat16_scalars_empty_and_noncontiguous_tensors(self):
+        scalar = torch.tensor(1.25, dtype=torch.bfloat16)
+        vector = scalar.reshape(1)
+        self.assertEqual(tensor_sha256(scalar), tensor_sha256(vector))
+
+        empty = torch.empty(0, dtype=torch.bfloat16)
+        self.assertEqual(tensor_sha256(empty), sha256_bytes(b""))
+
+        base = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+        noncontiguous = base.transpose(0, 1)
+        self.assertFalse(noncontiguous.is_contiguous())
+        self.assertEqual(
+            tensor_sha256(noncontiguous),
+            tensor_sha256(noncontiguous.contiguous()),
+        )
+
+    def test_artifact_round_trip_supports_bfloat16_scalar_tensor(self):
+        tensors = {
+            "adapter.scale": torch.tensor(1.0, dtype=torch.bfloat16),
+            "adapter.weight": torch.ones(2, 2, dtype=torch.bfloat16),
+        }
+        with tempfile.TemporaryDirectory() as temp_dir:
+            path, manifest = self._save(temp_dir, tensors=tensors)
+            loaded, loaded_manifest = load_artifact(path, expected_type="te_adapter")
+            self.assertEqual(manifest, loaded_manifest)
+            self.assertEqual(loaded["adapter.scale"].shape, torch.Size([]))
+            self.assertTrue(torch.equal(loaded["adapter.scale"], tensors["adapter.scale"]))
+
     def test_fingerprints_are_canonical_and_order_independent(self):
         self.assertEqual(config_fingerprint({"a": 1, "b": 2}), config_fingerprint({"b": 2, "a": 1}))
         self.assertNotEqual(config_fingerprint({"a": 1}), config_fingerprint({"a": 2}))

toolkit/trigger_binding_artifacts.py

@@ -129,7 +129,7 @@ def to_json_compatible(value: Any) -> Any:
         return {"__type__": "torch.dtype", "value": str(value).removeprefix("torch.")}
     if isinstance(value, torch.Tensor):
         tensor = value.detach().cpu().contiguous()
-        raw = tensor.view(torch.uint8).numpy().tobytes()
+        raw = _tensor_bytes(tensor)
         return {
             "__type__": "torch.Tensor",
             "dtype": str(tensor.dtype).removeprefix("torch."),
@@ -207,7 +207,13 @@ def decode_rng_state(encoded: Any) -> Any:
 
 
 def _tensor_bytes(tensor: torch.Tensor) -> bytes:
-    return tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()
+    normalized = tensor.detach().cpu().contiguous()
+    if normalized.numel() == 0:
+        return b""
+    # PyTorch cannot reinterpret a zero-dimensional tensor as a dtype with a
+    # different element size. Flattening preserves the exact storage bytes and
+    # also handles BF16, scalar adapter scales, and non-contiguous inputs.
+    return normalized.reshape(-1).view(torch.uint8).numpy().tobytes()
 
 
 def tensor_sha256(tensor: torch.Tensor) -> str: