Add SekoTalkInputs

Author: wangshankun

lightx2v/utils/input_info.py

@@ -86,58 +86,61 @@ class VaceInputInfo:
 
 @dataclass
 class S2VInputInfo:
-    infer_steps: int | Any = UNSET
-    seed: int | Any = UNSET
-    prompt: str | Any = UNSET
-    prompt_enhanced: str | Any = UNSET
-    negative_prompt: str | Any = UNSET
-    image_path: str | Any = UNSET
-    audio_path: str | Any = UNSET
-    audio_num: int | Any = UNSET
-    video_duration: float | Any = UNSET
-    with_mask: bool | Any = UNSET
-    return_result_tensor: bool | Any = UNSET
-    save_result_path: str | Any = UNSET
-    return_result_tensor: bool | Any = UNSET
-    stream_config: dict | Any = UNSET
-    resize_mode: str | Any = UNSET
-    target_shape: list | Any = UNSET
+    seed: int = field(default_factory=int)
+    prompt: str = field(default_factory=str)
+    prompt_enhanced: str = field(default_factory=str)
+    negative_prompt: str = field(default_factory=str)
+    image_path: str = field(default_factory=str)
+    audio_path: str = field(default_factory=str)
+    audio_num: int = field(default_factory=int)
+    with_mask: bool = field(default_factory=lambda: False)
+    save_result_path: str = field(default_factory=str)
+    return_result_tensor: bool = field(default_factory=lambda: False)
+    stream_config: dict = field(default_factory=dict)
+    # shape related
+    resize_mode: str = field(default_factory=str)
+    original_shape: list = field(default_factory=list)
+    resized_shape: list = field(default_factory=list)
+    latent_shape: list = field(default_factory=list)
+    target_shape: list = field(default_factory=list)
+
     # prev info
-    overlap_frame: torch.Tensor | Any = UNSET
-    overlap_latent: torch.Tensor | Any = UNSET
+    overlap_frame: torch.Tensor = field(default_factory=lambda: None)
+    overlap_latent: torch.Tensor = field(default_factory=lambda: None)
     # input preprocess audio
-    audio_clip: torch.Tensor | Any = UNSET
-
-    @classmethod
-    def from_args(cls, args, **overrides):
-        """
-        Build InputInfo from argparse.Namespace (or any object with __dict__)
-        Priority:
-            args < overrides
-        """
-        field_names = {f.name for f in fields(cls)}
-        data = {k: v for k, v in vars(args).items() if k in field_names}
-        data.update(overrides)
-        return cls(**data)
-
-    def normalize_unset_to_none(self):
-        """
-        Replace all UNSET fields with None.
-        Call this right before running / inference.
-        """
-        for f in fields(self):
-            if getattr(self, f.name) is UNSET:
-                setattr(self, f.name, None)
-        return self
+    audio_clip: torch.Tensor = field(default_factory=lambda: None)
 
 
 @dataclass
-class RS2VInputInfo(S2VInputInfo):
+class RS2VInputInfo:
+    seed: int = field(default_factory=int)
+    prompt: str = field(default_factory=str)
+    prompt_enhanced: str = field(default_factory=str)
+    negative_prompt: str = field(default_factory=str)
+    image_path: str = field(default_factory=str)
+    audio_path: str = field(default_factory=str)
+    audio_num: int = field(default_factory=int)
+    with_mask: bool = field(default_factory=lambda: False)
+    save_result_path: str = field(default_factory=str)
+    return_result_tensor: bool = field(default_factory=lambda: False)
+    stream_config: dict = field(default_factory=dict)
+    # shape related
+    resize_mode: str = field(default_factory=str)
+    original_shape: list = field(default_factory=list)
+    resized_shape: list = field(default_factory=list)
+    latent_shape: list = field(default_factory=list)
+    target_shape: list = field(default_factory=list)
+
+    # prev info
+    overlap_frame: torch.Tensor = field(default_factory=lambda: None)
+    overlap_latent: torch.Tensor = field(default_factory=lambda: None)
+    # input preprocess audio
+    audio_clip: torch.Tensor = field(default_factory=lambda: None)
     # input reference state
-    ref_state: int | Any = UNSET
+    ref_state: int = field(default_factory=int)
     # flags for first and last clip
-    is_first: bool | Any = UNSET
-    is_last: bool | Any = UNSET
+    is_first: bool = field(default_factory=lambda: False)
+    is_last: bool = field(default_factory=lambda: False)
 
 
 # Need Check
@@ -280,15 +283,6 @@ class WorldPlayT2VInputInfo:
     action: torch.Tensor = field(default_factory=lambda: None)
 
 
-def init_input_info_from_args(task, args, **overrides):
-    if task == "s2v":
-        return S2VInputInfo.from_args(args, **overrides)
-    elif task == "rs2v":
-        return RS2VInputInfo.from_args(args, **overrides)
-    else:
-        raise ValueError(f"Unsupported task: {task}")
-
-
 def init_empty_input_info(task):
     if task == "t2v":
         return T2VInputInfo()
@@ -320,6 +314,68 @@ def init_empty_input_info(task):
         raise ValueError(f"Unsupported task: {task}")
 
 
+@dataclass
+class SekoTalkInputs:
+    infer_steps: int | Any = UNSET
+    seed: int | Any = UNSET
+    prompt: str | Any = UNSET
+    prompt_enhanced: str | Any = UNSET
+    negative_prompt: str | Any = UNSET
+    image_path: str | Any = UNSET
+    audio_path: str | Any = UNSET
+    audio_num: int | Any = UNSET
+    video_duration: float | Any = UNSET
+    with_mask: bool | Any = UNSET
+    return_result_tensor: bool | Any = UNSET
+    save_result_path: str | Any = UNSET
+    return_result_tensor: bool | Any = UNSET
+    stream_config: dict | Any = UNSET
+
+    resize_mode: str | Any = UNSET
+    target_shape: list | Any = UNSET
+
+    # prev info
+    overlap_frame: torch.Tensor | Any = UNSET
+    overlap_latent: torch.Tensor | Any = UNSET
+    # input preprocess audio
+    audio_clip: torch.Tensor | Any = UNSET
+
+    # input reference state
+    ref_state: int | Any = UNSET
+    # flags for first and last clip
+    is_first: bool | Any = UNSET
+    is_last: bool | Any = UNSET
+
+    @classmethod
+    def from_args(cls, args, **overrides):
+        """
+        Build InputInfo from argparse.Namespace (or any object with __dict__)
+        Priority:
+            args < overrides
+        """
+        field_names = {f.name for f in fields(cls)}
+        data = {k: v for k, v in vars(args).items() if k in field_names}
+        data.update(overrides)
+        return cls(**data)
+
+    def normalize_unset_to_none(self):
+        """
+        Replace all UNSET fields with None.
+        Call this right before running / inference.
+        """
+        for f in fields(self):
+            if getattr(self, f.name) is UNSET:
+                setattr(self, f.name, None)
+        return self
+
+
+def init_input_info_from_args(task, args, **overrides):
+    if task in ["s2v", "rs2v"]:
+        return SekoTalkInputs.from_args(args, **overrides)
+    else:
+        raise ValueError(f"Unsupported task: {task}")
+
+
 def fill_input_info_from_defaults(input_info, defaults):
     for key in input_info.__dataclass_fields__:
         if key in defaults and getattr(input_info, key) is UNSET: