[Bugfix][sleep mode]: guard generation on partial wake and ensure wake idempotency (#4834)

Signed-off-by: vensen <vensenmu@gmail.com>

vllm_omni/diffusion/data.py

@@ -6,6 +6,7 @@ import os
 import random
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field, fields
+from enum import Enum
 from typing import TYPE_CHECKING, Any
 
 import diffusers
@@ -1488,5 +1489,12 @@ class OmniWakeTask:
     tags: list[str] | None = None
 
 
+class CuMemTag(str, Enum):
+    """Tags representing specific CuMem allocations for sleep/wake state tracking."""
+
+    WEIGHTS = "weights"
+    KV_CACHE = "kv_cache"
+
+
 # Special message broadcast via scheduler queues to signal worker shutdown.
 SHUTDOWN_MESSAGE = {"type": "shutdown"}