diff --git a/README.md b/README.md index 28ca11d..572c49b 100644 --- a/README.md +++ b/README.md @@ -64,7 +64,7 @@ loop.run_until_complete(main()) ```python from novelai_python.utils.random_prompt import RandomPromptGenerator -s = RandomPromptGenerator(nsfw_enabled=False).generate() +s = RandomPromptGenerator(nsfw_enabled=False).random_prompt() print(s) ``` diff --git a/playground/random_prompt.py b/playground/random_prompt.py index 1dd2981..40b124a 100644 --- a/playground/random_prompt.py +++ b/playground/random_prompt.py @@ -14,5 +14,5 @@ print(gen.get_weighted_choice([['m', 30], ['f', 50], ['o', 10]], ['m'])) print("====") for i in range(200): - s = RandomPromptGenerator(nsfw_enabled=True).generate() + s = RandomPromptGenerator(nsfw_enabled=True).random_prompt() print(s) diff --git a/src/novelai_python/utils/random_prompt/__init__.py b/src/novelai_python/utils/random_prompt/__init__.py index 3bf9db7..5de3a7d 100644 --- a/src/novelai_python/utils/random_prompt/__init__.py +++ b/src/novelai_python/utils/random_prompt/__init__.py @@ -199,8 +199,13 @@ def random_prompt(self, *, enable_moods: bool = True, enable_character: bool = True, enable_identity: bool = False, + must_appear=None, ): + if must_appear is None: + must_appear = [] tags = [] + # 必须出现的标签 + tags.extend(must_appear) if self.nsfw_enabled: tags.append('nsfw') if random.random() < 0.1: @@ -208,7 +213,7 @@ def random_prompt(self, *, tags.append('lewd') irs = self.get_weighted_choice([[1, 70], [2, 20], [3, 7], [0, 5]], tags) if self.nsfw_enabled: - irs = self.get_weighted_choice([[1, 35], [2, 20], [3, 7]], tags) + irs = self.get_weighted_choice([[1, 40], [2, 20], [3, 7]], tags) if irs == 0: tags.append('no humans') if random.random() < 0.3: