Class DashScopePromptExpander

wan/utils/prompt_extend.py:213–361

class DashScopePromptExpander(PromptExpander):

    def __init__(self,
                 api_key=None,
                 model_name=None,
                 max_image_size=512 * 512,
                 retry_times=4,
                 is_vl=False,
                 **kwargs):
        '''
        Args:
            api_key: The API key for Dash Scope authentication and access to related services.
            model_name: Model name, 'qwen-plus' for extending prompts, 'qwen-vl-max' for extending prompt-images.
            max_image_size: The maximum size of the image; unit unspecified (e.g., pixels, KB). Please specify the unit based on actual usage.
            retry_times: Number of retry attempts in case of request failure.
            is_vl: A flag indicating whether the task involves visual-language processing.
            **kwargs: Additional keyword arguments that can be passed to the function or method.
        '''
        if model_name is None:
            model_name = 'qwen-plus' if not is_vl else 'qwen-vl-max'
        super().__init__(model_name, is_vl, **kwargs)
        if api_key is not None:
            dashscope.api_key = api_key
        elif 'DASH_API_KEY' in os.environ and os.environ[
                'DASH_API_KEY'] is not None:
            dashscope.api_key = os.environ['DASH_API_KEY']
        else:
            raise ValueError("DASH_API_KEY is not set")
        if 'DASH_API_URL' in os.environ and os.environ[
                'DASH_API_URL'] is not None:
            dashscope.base_http_api_url = os.environ['DASH_API_URL']
        else:
            dashscope.base_http_api_url = 'https://dashscope.aliyuncs.com/api/v1'
        self.api_key = api_key

        self.max_image_size = max_image_size
        self.model = model_name
        self.retry_times = retry_times

    def extend(self, prompt, system_prompt, seed=-1, *args, **kwargs):
        messages = [{
            'role': 'system',
            'content': system_prompt
        }, {
            'role': 'user',
            'content': prompt
        }]

        exception = None
        for _ in range(self.retry_times):
            try:
                response = dashscope.Generation.call(
                    self.model,
                    messages=messages,
                    seed=seed,
                    result_format='message',  # set the result to be "message" format.
                )
                assert response.status_code == HTTPStatus.OK, response