fix: adapt model_manager to HF model API (no schedule_config/cfg_config dicts)

model_manager.py

@@ -62,9 +62,14 @@ class ModelManager:
         # Load models from HF Hub
         self.vae, self.model = self._load_models(model_name)
 
-        # Save clean copies of user-facing configs (before any runtime injection)
-        self._base_schedule_config = dict(self.model.schedule_config)
-        self._base_cfg_config = dict(self.model.cfg_config)
+        # Build config dicts from model's individual attributes (HF model API)
+        self._base_schedule_config = {
+            "chunk_size": self.model.chunk_size,
+            "steps": self.model.noise_steps,
+        }
+        self._base_cfg_config = {
+            "cfg_scale": self.model.cfg_scale,
+        }
 
         # Frame buffer
         self.frame_buffer = FrameBuffer(target_buffer_size=4)
@@ -82,7 +87,8 @@ class ModelManager:
         self.should_stop = False
 
         # Model generation state
-        self.first_chunk = True
+        self.first_chunk = True  # For VAE stream_decode
+        self._model_first_chunk = True  # For model stream_generate_step
         self.history_length = 30
 
         print("ModelManager initialized successfully")
@@ -187,16 +193,14 @@ class ModelManager:
             self.stream_recovery.reset()
             self.vae.clear_cache()
             self.first_chunk = True
-            # Restore clean config before init (clears runtime-injected keys)
-            self.model.schedule_config.clear()
-            self.model.schedule_config.update(self._base_schedule_config)
-            self.model.init_generated(
-                self.history_length,
-                batch_size=1,
-                schedule_config=self.model.schedule_config,
-            )
+            self._model_first_chunk = True
+            # Restore model params from base config
+            self.model.chunk_size = self._base_schedule_config["chunk_size"]
+            self.model.noise_steps = self._base_schedule_config["steps"]
+            self.model.cfg_scale = self._base_cfg_config["cfg_scale"]
+            self.model.init_generated(self.history_length, batch_size=1)
             print(
-                f"Model initialized with history length: {self.history_length}, schedule_config: {self.model.schedule_config}"
+                f"Model initialized with history length: {self.history_length}"
             )
 
             # Start generation thread
@@ -269,20 +273,16 @@ class ModelManager:
             joints_num=22, smoothing_alpha=self.smoothing_alpha
         )
 
-        # Restore clean configs before init (clears runtime-injected keys)
-        self.model.schedule_config.clear()
-        self.model.schedule_config.update(self._base_schedule_config)
-        self.model.cfg_config.clear()
-        self.model.cfg_config.update(self._base_cfg_config)
-
-        # Initialize model (reads steps/chunk_size from model.schedule_config directly)
-        self.model.init_generated(
-            self.history_length,
-            batch_size=1,
-            schedule_config=self.model.schedule_config,
-        )
+        # Restore model params from base config
+        self.model.chunk_size = self._base_schedule_config["chunk_size"]
+        self.model.noise_steps = self._base_schedule_config["steps"]
+        self.model.cfg_scale = self._base_cfg_config["cfg_scale"]
+        self._model_first_chunk = True
+
+        # Initialize model
+        self.model.init_generated(self.history_length, batch_size=1)
         print(
-            f"Model reset - history: {self.history_length}, smoothing: {self.smoothing_alpha}, schedule_config: {self.model.schedule_config}"
+            f"Model reset - history: {self.history_length}, smoothing: {self.smoothing_alpha}"
         )
 
     def _generation_loop(self):
@@ -299,12 +299,14 @@ class ModelManager:
                     try:
                         step_start = time.time()
 
-                        # Generate one token (produces 4 frames from VAE)
-                        text_key = self.model.input_keys["text"]
-                        x = {text_key: [self.current_text]}
+                        # Generate one token (produces frames from VAE)
+                        x = {"text": [self.current_text]}
 
                         # Generate from model (1 token)
-                        output = self.model.stream_generate_step(x)
+                        output = self.model.stream_generate_step(
+                            x, first_chunk=self._model_first_chunk
+                        )
+                        self._model_first_chunk = False
                         generated = output["generated"]
 
                         # Skip if no frames committed yet
@@ -364,8 +366,13 @@ class ModelManager:
             "current_text": self.current_text,
             "smoothing_alpha": self.smoothing_alpha,
             "history_length": self.history_length,
-            "schedule_config": dict(self.model.schedule_config),
-            "cfg_config": dict(self.model.cfg_config),
+            "schedule_config": {
+                "chunk_size": self.model.chunk_size,
+                "steps": self.model.noise_steps,
+            },
+            "cfg_config": {
+                "cfg_scale": self.model.cfg_scale,
+            },
         }