support wanx prompter

This commit is contained in:
mi804
2025-02-20 16:08:22 +08:00
parent 427232cbc0
commit 0c92f3b2cc
4 changed files with 376 additions and 0 deletions

View File

@@ -0,0 +1,18 @@
import torch
from diffsynth.prompters import WanXPrompter
from diffsynth.models.wanx_text_encoder import WanXTextEncoder
prompter = WanXPrompter('models/WanX/google/umt5-xxl')
text_encoder = WanXTextEncoder()
text_encoder.load_state_dict(torch.load('models/WanX/models_t5_umt5-xxl-enc-bf16.pth', map_location='cpu'))
text_encoder = text_encoder.eval().requires_grad_(False).to(dtype=torch.bfloat16, device='cuda')
prompter.fetch_models(text_encoder)
prompt = '维京战士双手挥舞着大斧,对抗猛犸象,黄昏,雪地中,漫天飞雪'
neg_prompt = '色调艳丽过曝静态细节模糊不清字幕风格作品画作画面静止整体发灰最差质量低质量JPEG压缩残留丑陋的残缺的多余的手指画得不好的手部画得不好的脸部畸形的毁容的形态畸形的肢体手指融合静止不动的画面杂乱的背景三条腿背景人很多倒着走'
prompt_emb = prompter.encode_prompt(prompt)
neg_prompt_emb = prompter.encode_prompt(neg_prompt)
print(prompt_emb[0]) # torch.Size([31, 4096])
print(neg_prompt_emb[0]) # torch.Size([126, 4096])