wehub-resource-sync
543 lines
20 KiB
Python
543 lines
20 KiB
Python
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import random
|
|
import sys
|
|
import tempfile
|
|
from dataclasses import dataclass
|
|
from http import HTTPStatus
|
|
from typing import Optional, Union
|
|
|
|
import dashscope
|
|
import torch
|
|
from PIL import Image
|
|
|
|
try:
|
|
from flash_attn import flash_attn_varlen_func
|
|
FLASH_VER = 2
|
|
except ModuleNotFoundError:
|
|
flash_attn_varlen_func = None # in compatible with CPU machines
|
|
FLASH_VER = None
|
|
|
|
from .system_prompt import *
|
|
|
|
DEFAULT_SYS_PROMPTS = {
|
|
"t2v-A14B": {
|
|
"zh": T2V_A14B_ZH_SYS_PROMPT,
|
|
"en": T2V_A14B_EN_SYS_PROMPT,
|
|
},
|
|
"i2v-A14B": {
|
|
"zh": I2V_A14B_ZH_SYS_PROMPT,
|
|
"en": I2V_A14B_EN_SYS_PROMPT,
|
|
"empty": {
|
|
"zh": I2V_A14B_EMPTY_ZH_SYS_PROMPT,
|
|
"en": I2V_A14B_EMPTY_EN_SYS_PROMPT,
|
|
}
|
|
},
|
|
"ti2v-5B": {
|
|
"t2v": {
|
|
"zh": T2V_A14B_ZH_SYS_PROMPT,
|
|
"en": T2V_A14B_EN_SYS_PROMPT,
|
|
},
|
|
"i2v": {
|
|
"zh": I2V_A14B_ZH_SYS_PROMPT,
|
|
"en": I2V_A14B_EN_SYS_PROMPT,
|
|
}
|
|
},
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class PromptOutput(object):
|
|
status: bool
|
|
prompt: str
|
|
seed: int
|
|
system_prompt: str
|
|
message: str
|
|
|
|
def add_custom_field(self, key: str, value) -> None:
|
|
self.__setattr__(key, value)
|
|
|
|
|
|
class PromptExpander:
|
|
|
|
def __init__(self, model_name, task, is_vl=False, device=0, **kwargs):
|
|
self.model_name = model_name
|
|
self.task = task
|
|
self.is_vl = is_vl
|
|
self.device = device
|
|
|
|
def extend_with_img(self,
|
|
prompt,
|
|
system_prompt,
|
|
image=None,
|
|
seed=-1,
|
|
*args,
|
|
**kwargs):
|
|
pass
|
|
|
|
def extend(self, prompt, system_prompt, seed=-1, *args, **kwargs):
|
|
pass
|
|
|
|
def decide_system_prompt(self, tar_lang="zh", prompt=None):
|
|
assert self.task is not None
|
|
if "ti2v" in self.task:
|
|
if self.is_vl:
|
|
return DEFAULT_SYS_PROMPTS[self.task]["i2v"][tar_lang]
|
|
else:
|
|
return DEFAULT_SYS_PROMPTS[self.task]["t2v"][tar_lang]
|
|
if "i2v" in self.task and len(prompt) == 0:
|
|
return DEFAULT_SYS_PROMPTS[self.task]["empty"][tar_lang]
|
|
return DEFAULT_SYS_PROMPTS[self.task][tar_lang]
|
|
|
|
def __call__(self,
|
|
prompt,
|
|
system_prompt=None,
|
|
tar_lang="zh",
|
|
image=None,
|
|
seed=-1,
|
|
*args,
|
|
**kwargs):
|
|
if system_prompt is None:
|
|
system_prompt = self.decide_system_prompt(
|
|
tar_lang=tar_lang, prompt=prompt)
|
|
if seed < 0:
|
|
seed = random.randint(0, sys.maxsize)
|
|
if image is not None and self.is_vl:
|
|
return self.extend_with_img(
|
|
prompt, system_prompt, image=image, seed=seed, *args, **kwargs)
|
|
elif not self.is_vl:
|
|
return self.extend(prompt, system_prompt, seed, *args, **kwargs)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
|
|
class DashScopePromptExpander(PromptExpander):
|
|
|
|
def __init__(self,
|
|
api_key=None,
|
|
model_name=None,
|
|
task=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.
|
|
task: Task name. This is required to determine the default system prompt.
|
|
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, task, 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
|
|
expanded_prompt = response['output']['choices'][0]['message'][
|
|
'content']
|
|
return PromptOutput(
|
|
status=True,
|
|
prompt=expanded_prompt,
|
|
seed=seed,
|
|
system_prompt=system_prompt,
|
|
message=json.dumps(response, ensure_ascii=False))
|
|
except Exception as e:
|
|
exception = e
|
|
return PromptOutput(
|
|
status=False,
|
|
prompt=prompt,
|
|
seed=seed,
|
|
system_prompt=system_prompt,
|
|
message=str(exception))
|
|
|
|
def extend_with_img(self,
|
|
prompt,
|
|
system_prompt,
|
|
image: Union[Image.Image, str] = None,
|
|
seed=-1,
|
|
*args,
|
|
**kwargs):
|
|
if isinstance(image, str):
|
|
image = Image.open(image).convert('RGB')
|
|
w = image.width
|
|
h = image.height
|
|
area = min(w * h, self.max_image_size)
|
|
aspect_ratio = h / w
|
|
resized_h = round(math.sqrt(area * aspect_ratio))
|
|
resized_w = round(math.sqrt(area / aspect_ratio))
|
|
image = image.resize((resized_w, resized_h))
|
|
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
|
|
image.save(f.name)
|
|
fname = f.name
|
|
image_path = f"file://{f.name}"
|
|
prompt = f"{prompt}"
|
|
messages = [
|
|
{
|
|
'role': 'system',
|
|
'content': [{
|
|
"text": system_prompt
|
|
}]
|
|
},
|
|
{
|
|
'role': 'user',
|
|
'content': [{
|
|
"text": prompt
|
|
}, {
|
|
"image": image_path
|
|
}]
|
|
},
|
|
]
|
|
response = None
|
|
result_prompt = prompt
|
|
exception = None
|
|
status = False
|
|
for _ in range(self.retry_times):
|
|
try:
|
|
response = dashscope.MultiModalConversation.call(
|
|
self.model,
|
|
messages=messages,
|
|
seed=seed,
|
|
result_format='message', # set the result to be "message" format.
|
|
)
|
|
assert response.status_code == HTTPStatus.OK, response
|
|
result_prompt = response['output']['choices'][0]['message'][
|
|
'content'][0]['text'].replace('\n', '\\n')
|
|
status = True
|
|
break
|
|
except Exception as e:
|
|
exception = e
|
|
result_prompt = result_prompt.replace('\n', '\\n')
|
|
os.remove(fname)
|
|
|
|
return PromptOutput(
|
|
status=status,
|
|
prompt=result_prompt,
|
|
seed=seed,
|
|
system_prompt=system_prompt,
|
|
message=str(exception) if not status else json.dumps(
|
|
response, ensure_ascii=False))
|
|
|
|
|
|
class QwenPromptExpander(PromptExpander):
|
|
model_dict = {
|
|
"QwenVL2.5_3B": "Qwen/Qwen2.5-VL-3B-Instruct",
|
|
"QwenVL2.5_7B": "Qwen/Qwen2.5-VL-7B-Instruct",
|
|
"Qwen2.5_3B": "Qwen/Qwen2.5-3B-Instruct",
|
|
"Qwen2.5_7B": "Qwen/Qwen2.5-7B-Instruct",
|
|
"Qwen2.5_14B": "Qwen/Qwen2.5-14B-Instruct",
|
|
}
|
|
|
|
def __init__(self,
|
|
model_name=None,
|
|
task=None,
|
|
device=0,
|
|
is_vl=False,
|
|
**kwargs):
|
|
'''
|
|
Args:
|
|
model_name: Use predefined model names such as 'QwenVL2.5_7B' and 'Qwen2.5_14B',
|
|
which are specific versions of the Qwen model. Alternatively, you can use the
|
|
local path to a downloaded model or the model name from Hugging Face."
|
|
Detailed Breakdown:
|
|
Predefined Model Names:
|
|
* 'QwenVL2.5_7B' and 'Qwen2.5_14B' are specific versions of the Qwen model.
|
|
Local Path:
|
|
* You can provide the path to a model that you have downloaded locally.
|
|
Hugging Face Model Name:
|
|
* You can also specify the model name from Hugging Face's model hub.
|
|
task: Task name. This is required to determine the default system prompt.
|
|
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 = 'Qwen2.5_14B' if not is_vl else 'QwenVL2.5_7B'
|
|
super().__init__(model_name, task, is_vl, device, **kwargs)
|
|
if (not os.path.exists(self.model_name)) and (self.model_name
|
|
in self.model_dict):
|
|
self.model_name = self.model_dict[self.model_name]
|
|
|
|
if self.is_vl:
|
|
# default: Load the model on the available device(s)
|
|
from transformers import (
|
|
AutoProcessor,
|
|
AutoTokenizer,
|
|
Qwen2_5_VLForConditionalGeneration,
|
|
)
|
|
try:
|
|
from .qwen_vl_utils import process_vision_info
|
|
except:
|
|
from qwen_vl_utils import process_vision_info
|
|
self.process_vision_info = process_vision_info
|
|
min_pixels = 256 * 28 * 28
|
|
max_pixels = 1280 * 28 * 28
|
|
self.processor = AutoProcessor.from_pretrained(
|
|
self.model_name,
|
|
min_pixels=min_pixels,
|
|
max_pixels=max_pixels,
|
|
use_fast=True)
|
|
self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
|
self.model_name,
|
|
torch_dtype=torch.bfloat16 if FLASH_VER == 2 else
|
|
torch.float16 if "AWQ" in self.model_name else "auto",
|
|
attn_implementation="flash_attention_2"
|
|
if FLASH_VER == 2 else None,
|
|
device_map="cpu")
|
|
else:
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
self.model = AutoModelForCausalLM.from_pretrained(
|
|
self.model_name,
|
|
torch_dtype=torch.float16
|
|
if "AWQ" in self.model_name else "auto",
|
|
attn_implementation="flash_attention_2"
|
|
if FLASH_VER == 2 else None,
|
|
device_map="cpu")
|
|
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
|
|
|
|
def extend(self, prompt, system_prompt, seed=-1, *args, **kwargs):
|
|
self.model = self.model.to(self.device)
|
|
messages = [{
|
|
"role": "system",
|
|
"content": system_prompt
|
|
}, {
|
|
"role": "user",
|
|
"content": prompt
|
|
}]
|
|
text = self.tokenizer.apply_chat_template(
|
|
messages, tokenize=False, add_generation_prompt=True)
|
|
model_inputs = self.tokenizer([text],
|
|
return_tensors="pt").to(self.model.device)
|
|
|
|
generated_ids = self.model.generate(**model_inputs, max_new_tokens=512)
|
|
generated_ids = [
|
|
output_ids[len(input_ids):] for input_ids, output_ids in zip(
|
|
model_inputs.input_ids, generated_ids)
|
|
]
|
|
|
|
expanded_prompt = self.tokenizer.batch_decode(
|
|
generated_ids, skip_special_tokens=True)[0]
|
|
self.model = self.model.to("cpu")
|
|
return PromptOutput(
|
|
status=True,
|
|
prompt=expanded_prompt,
|
|
seed=seed,
|
|
system_prompt=system_prompt,
|
|
message=json.dumps({"content": expanded_prompt},
|
|
ensure_ascii=False))
|
|
|
|
def extend_with_img(self,
|
|
prompt,
|
|
system_prompt,
|
|
image: Union[Image.Image, str] = None,
|
|
seed=-1,
|
|
*args,
|
|
**kwargs):
|
|
self.model = self.model.to(self.device)
|
|
messages = [{
|
|
'role': 'system',
|
|
'content': [{
|
|
"type": "text",
|
|
"text": system_prompt
|
|
}]
|
|
}, {
|
|
"role":
|
|
"user",
|
|
"content": [
|
|
{
|
|
"type": "image",
|
|
"image": image,
|
|
},
|
|
{
|
|
"type": "text",
|
|
"text": prompt
|
|
},
|
|
],
|
|
}]
|
|
|
|
# Preparation for inference
|
|
text = self.processor.apply_chat_template(
|
|
messages, tokenize=False, add_generation_prompt=True)
|
|
image_inputs, video_inputs = self.process_vision_info(messages)
|
|
inputs = self.processor(
|
|
text=[text],
|
|
images=image_inputs,
|
|
videos=video_inputs,
|
|
padding=True,
|
|
return_tensors="pt",
|
|
)
|
|
inputs = inputs.to(self.device)
|
|
|
|
# Inference: Generation of the output
|
|
generated_ids = self.model.generate(**inputs, max_new_tokens=512)
|
|
generated_ids_trimmed = [
|
|
out_ids[len(in_ids):]
|
|
for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
|
|
]
|
|
expanded_prompt = self.processor.batch_decode(
|
|
generated_ids_trimmed,
|
|
skip_special_tokens=True,
|
|
clean_up_tokenization_spaces=False)[0]
|
|
self.model = self.model.to("cpu")
|
|
return PromptOutput(
|
|
status=True,
|
|
prompt=expanded_prompt,
|
|
seed=seed,
|
|
system_prompt=system_prompt,
|
|
message=json.dumps({"content": expanded_prompt},
|
|
ensure_ascii=False))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="[%(asctime)s] %(levelname)s: %(message)s",
|
|
handlers=[logging.StreamHandler(stream=sys.stdout)])
|
|
|
|
seed = 100
|
|
prompt = "夏日海滩度假风格，一只戴着墨镜的白色猫咪坐在冲浪板上。猫咪毛发蓬松，表情悠闲，直视镜头。背景是模糊的海滩景色，海水清澈，远处有绿色的山丘和蓝天白云。猫咪的姿态自然放松，仿佛在享受海风和阳光。近景特写，强调猫咪的细节和海滩的清新氛围。"
|
|
en_prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
|
|
image = "./examples/i2v_input.JPG"
|
|
|
|
def test(method,
|
|
prompt,
|
|
model_name,
|
|
task,
|
|
image=None,
|
|
en_prompt=None,
|
|
seed=None):
|
|
prompt_expander = method(
|
|
model_name=model_name, task=task, is_vl=image is not None)
|
|
result = prompt_expander(prompt, image=image, tar_lang="zh")
|
|
logging.info(f"zh prompt -> zh: {result.prompt}")
|
|
result = prompt_expander(prompt, image=image, tar_lang="en")
|
|
logging.info(f"zh prompt -> en: {result.prompt}")
|
|
if en_prompt is not None:
|
|
result = prompt_expander(en_prompt, image=image, tar_lang="zh")
|
|
logging.info(f"en prompt -> zh: {result.prompt}")
|
|
result = prompt_expander(en_prompt, image=image, tar_lang="en")
|
|
logging.info(f"en prompt -> en: {result.prompt}")
|
|
|
|
ds_model_name = None
|
|
ds_vl_model_name = None
|
|
qwen_model_name = None
|
|
qwen_vl_model_name = None
|
|
|
|
for task in ["t2v-A14B", "i2v-A14B", "ti2v-5B"]:
|
|
# test prompt extend
|
|
if "t2v" in task or "ti2v" in task:
|
|
# test dashscope api
|
|
logging.info(f"-" * 40)
|
|
logging.info(f"Testing {task} dashscope prompt extend")
|
|
test(
|
|
DashScopePromptExpander,
|
|
prompt,
|
|
ds_model_name,
|
|
task,
|
|
image=None,
|
|
en_prompt=en_prompt,
|
|
seed=seed)
|
|
|
|
# test qwen api
|
|
logging.info(f"-" * 40)
|
|
logging.info(f"Testing {task} qwen prompt extend")
|
|
test(
|
|
QwenPromptExpander,
|
|
prompt,
|
|
qwen_model_name,
|
|
task,
|
|
image=None,
|
|
en_prompt=en_prompt,
|
|
seed=seed)
|
|
|
|
# test prompt-image extend
|
|
if "i2v" in task:
|
|
# test dashscope api
|
|
logging.info(f"-" * 40)
|
|
logging.info(f"Testing {task} dashscope vl prompt extend")
|
|
test(
|
|
DashScopePromptExpander,
|
|
prompt,
|
|
ds_vl_model_name,
|
|
task,
|
|
image=image,
|
|
en_prompt=en_prompt,
|
|
seed=seed)
|
|
|
|
# test qwen api
|
|
logging.info(f"-" * 40)
|
|
logging.info(f"Testing {task} qwen vl prompt extend")
|
|
test(
|
|
QwenPromptExpander,
|
|
prompt,
|
|
qwen_vl_model_name,
|
|
task,
|
|
image=image,
|
|
en_prompt=en_prompt,
|
|
seed=seed)
|
|
|
|
# test empty prompt extend
|
|
if "i2v-A14B" in task:
|
|
# test dashscope api
|
|
logging.info(f"-" * 40)
|
|
logging.info(f"Testing {task} dashscope vl empty prompt extend")
|
|
test(
|
|
DashScopePromptExpander,
|
|
"",
|
|
ds_vl_model_name,
|
|
task,
|
|
image=image,
|
|
en_prompt=None,
|
|
seed=seed)
|
|
|
|
# test qwen api
|
|
logging.info(f"-" * 40)
|
|
logging.info(f"Testing {task} qwen vl empty prompt extend")
|
|
test(
|
|
QwenPromptExpander,
|
|
"",
|
|
qwen_vl_model_name,
|
|
task,
|
|
image=image,
|
|
en_prompt=None,
|
|
seed=seed)
|