Bladeren bron

feat:添加generate_long_parallel_batched

zhaohaipeng 2 maanden geleden
bovenliggende
commit
85c3e9d04d
1 gewijzigde bestanden met toevoegingen van 14 en 10 verwijderingen
  1. 14 10
      fish_speech/models/text2semantic/inference.py

+ 14 - 10
fish_speech/models/text2semantic/inference.py

@@ -998,16 +998,20 @@ def generate_batched(
 def generate_long_parallel_batched(
         *,
         model,
-        device,
-        decode_one_token,
-        text,
-        prompt_text=None,
-        prompt_tokens=None,
-        max_new_tokens=512,
-        chunk_length=512,
-        temperature=1.0,
-        top_p=0.9,
-        top_k=30,
+        device: Union[str, torch.device],
+        decode_one_token: Callable,
+        text: str,
+        num_samples: int = 1,
+        max_new_tokens: int = 0,
+        top_p: float = 0.9,
+        top_k: int = 30,
+        repetition_penalty: float = 1.1,
+        temperature: float = 1.0,
+        compile: bool = False,
+        iterative_prompt: bool = True,
+        chunk_length: int = 512,
+        prompt_text: Optional[Union[str, list[str]]] = None,
+        prompt_tokens: Optional[Union[torch.Tensor, list[torch.Tensor]]] = None,
 ):
     """
     最小侵入版本: