37 lines
1.6 KiB
Python
37 lines
1.6 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import torch
|
|
|
|
class PromptExtendAnnotator:
|
|
def __init__(self, cfg, device=None):
|
|
from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander
|
|
self.mode = cfg.get('MODE', "local_qwen")
|
|
self.model_name = cfg.get('MODEL_NAME', "Qwen2.5_3B")
|
|
self.is_vl = cfg.get('IS_VL', False)
|
|
self.system_prompt = cfg.get('SYSTEM_PROMPT', None)
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else device
|
|
self.device_id = self.device.index if self.device.type == 'cuda' else None
|
|
rank = self.device_id if self.device_id is not None else 0
|
|
if self.mode == "dashscope":
|
|
self.prompt_expander = DashScopePromptExpander(
|
|
model_name=self.model_name, is_vl=self.is_vl)
|
|
elif self.mode == "local_qwen":
|
|
self.prompt_expander = QwenPromptExpander(
|
|
model_name=self.model_name,
|
|
is_vl=self.is_vl,
|
|
device=rank)
|
|
else:
|
|
raise NotImplementedError(f"Unsupport prompt_extend_method: {self.mode}")
|
|
|
|
|
|
def forward(self, prompt, system_prompt=None, seed=-1):
|
|
system_prompt = system_prompt if system_prompt is not None else self.system_prompt
|
|
output = self.prompt_expander(prompt, system_prompt=system_prompt, seed=seed)
|
|
if output.status == False:
|
|
print(f"Extending prompt failed: {output.message}")
|
|
output_prompt = prompt
|
|
else:
|
|
output_prompt = output.prompt
|
|
return output_prompt
|