[diffusion] refactor: unify model loading and offloading behavior (#15923)

2025-12-27 16:18:24 +08:00
parent 171912a9e3
commit aa89c6a7e2
9 changed files with 97 additions and 38 deletions
@@ -108,11 +108,11 @@ class GPUWorker:
                peak_memory_bytes = torch.cuda.max_memory_allocated()
                output_batch.peak_memory_mb = peak_memory_bytes / (1024**2)

-            if output_batch.timings:
-                duration_ms = (time.monotonic() - start_time) * 1000
-                output_batch.timings.total_duration_ms = duration_ms
-                if StageProfiler.metrics_enabled():
-                    PerformanceLogger.log_request_summary(timings=output_batch.timings)
+            duration_ms = (time.monotonic() - start_time) * 1000
+            output_batch.timings.total_duration_ms = duration_ms
+
+            if StageProfiler.metrics_enabled():
+                PerformanceLogger.log_request_summary(timings=output_batch.timings)
        except Exception as e:
            logger.error(
                f"Error executing request {req.request_id}: {e}", exc_info=True