206 lines
11 KiB
Python
206 lines
11 KiB
Python
from transformers import MarianMTModel, MarianTokenizer
|
||
import re
|
||
import json
|
||
import requests
|
||
from ollama import Client, Options
|
||
|
||
# 去这个地址： https://huggingface.co/Helsinki-NLP 可以找到更多语种的模型，用同样的格式添加到如下列表，以便下次选用，翻译模型会被自动下载至本地
|
||
marian_list = [
|
||
"opus-mt-zh-en",
|
||
"opus-mt-ru-en",
|
||
"opus-mt-th-en",
|
||
"opus-mt-es-en",
|
||
"opus-mt-fr-en",
|
||
"opus-mt-ar-en",
|
||
"opus-mt-hi-en",
|
||
"opus-mt-bn-en",
|
||
"opus-mt-pt-en",
|
||
"opus-mt-ja-en",
|
||
"opus-mt-de-en",
|
||
"opus-mt-ko-en",
|
||
"opus-mt-vi-en",
|
||
"opus-mt-tr-en",
|
||
"opus-mt-nl-en",
|
||
]
|
||
|
||
class MTCLIPEncode:
|
||
def __init__(self):
|
||
self.default_ollama_url = "http://127.0.0.1:11434"
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"clip": ("CLIP", {}),
|
||
"checkpoint": (marian_list, {"multiline": False, "default": "opus-mt-zh-en"}),
|
||
"text": ("STRING", {"multiline": True, "default": " prefix | translate_part | suffix"}),
|
||
"ollama_url": ("STRING", {"default": "http://127.0.0.1:11434"}),
|
||
"ollama_model": ("STRING", {"default": "mistral-nemo:latest"}),
|
||
}
|
||
}
|
||
|
||
# 确保双竖线 `||` 中间的内容为汉字或不少于三个字母的非英文单词
|
||
def is_valid_translate_part(self, text):
|
||
return bool(re.search(r'[\u4e00-\u9fff]|\b\w{3,}\b', text))
|
||
|
||
def optimize_with_ollama(self, text, url, model):
|
||
client = Client(host=url)
|
||
|
||
prompt = f"""
|
||
You are a talented art director with a deep understanding of visual arts. You excel at transforming simple text descriptions into vivid, imaginative images. Now, you will be playing the role of a prompt engineer, providing precise and creative prompts for the Stable Diffusion model.
|
||
|
||
Please convert the following user input text: {text}, into a detailed and vivid prompt to inspire Stable Diffusion to generate a captivating image.
|
||
|
||
**Focus on:**
|
||
* **Clarity and precision:** Use clear and concise language that Stable Diffusion can easily understand.
|
||
* **Visual storytelling:** Create compelling narratives through your descriptions.
|
||
* **Artistic techniques:** Employ a variety of artistic techniques like lighting, composition, and perspective.
|
||
* **Style and atmosphere:** Convey the desired mood and style using vivid language and metaphors.
|
||
|
||
**Consider incorporating:**
|
||
* **Specific art styles:** e.g., impressionism, surrealism, cyberpunk
|
||
* **Cinematographic techniques:** e.g., shallow depth of field, wide-angle shot
|
||
* **Descriptive adjectives:** e.g., ethereal, luminous, brooding
|
||
* **Compositional elements:** e.g., rule of thirds, golden ratio
|
||
|
||
**Example:**
|
||
* **User input:** A lonely robot sitting on a distant planet
|
||
* **Your prompt:** # solitary android, desolate alien planet, distant star, panoramic view, sci-fi aesthetic, melancholic atmosphere, cinematic lighting #
|
||
|
||
**Remember:** Your goal is to create a prompt that inspires Stable Diffusion to generate a visually stunning and emotionally resonant image.
|
||
|
||
**Please enclose your prompt within two `#` symbols.**
|
||
"""
|
||
|
||
response = client.generate(model=model, prompt=prompt)
|
||
return response['response']
|
||
|
||
def mtencode(self, clip, checkpoint, text, ollama_url, ollama_model):
|
||
|
||
print(f" 🫐 🫐 🫐 🫐 🫐 🫐") # 打印分隔线
|
||
|
||
pattern = r"([^|]*)\|([^|]*)\|([^|]*)" # 需要调用翻译模型的提示词基本格式
|
||
match = re.match(pattern, text) # 判断是否需要调用翻译模型
|
||
|
||
# 若输入 text 中，没有`||`(没有翻译请求)，判断叹号的数量和位置
|
||
ollama_for_original = text.count('!') == 1 and '|' not in text # 对应的输入格式为："English words English words ! English words"，叹号 `!` 为一个，可以在任意位置
|
||
ollama_for_part_of_original = text.count('!') == 2 and '|' not in text # 对应格式为："English prefix ! English words for Ollama refine ! English suffix"
|
||
|
||
# 有`||`，说明有翻译请求，进行翻译处理
|
||
if match:
|
||
prefix, translate_part, suffix = [re.sub(r'^[\s|,]+|[\s|,]+$', '', part.strip()) for part in match.groups()] # 被 `||` 双竖线分隔的三个部分
|
||
print(f" \033[94m{prefix}\033[0m, \033[92m{translate_part}\033[0m, \033[94m{suffix}\033[0m") # 打印用户的原始输入，并用颜色区分三个部分
|
||
|
||
# ollama_in_translate_part 和 ollama_for_all_3_parts 的判断在解析出 prefix、translate_part 和 suffix 之后
|
||
ollama_in_translate_part = '!' in translate_part # 判断叹号是否在翻译部分
|
||
ollama_for_all_3_parts = (
|
||
not ollama_in_translate_part and (
|
||
'!' in prefix or '!' in suffix # 判断叹号是否在前缀或后缀中
|
||
)
|
||
)
|
||
|
||
if not self.is_valid_translate_part(translate_part):
|
||
prompt_text = f"{prefix}, {suffix}"
|
||
else:
|
||
model_name = 'Helsinki-NLP/' + checkpoint
|
||
tokenizer = MarianTokenizer.from_pretrained(model_name)
|
||
model = MarianMTModel.from_pretrained(model_name)
|
||
|
||
if ollama_in_translate_part:
|
||
|
||
translate_part = translate_part.replace('!', '') # 去掉 `!`
|
||
translated = model.generate(**tokenizer(translate_part, return_tensors="pt", padding=True))
|
||
translated_text = re.sub(r'^[\s,.]+|[\s,.]+$', '', tokenizer.decode(translated[0], skip_special_tokens=True))
|
||
print(f" \033[94m{prefix}\033[0m, \033[92m{translated_text}\033[0m, \033[94m{suffix}\033[0m") # 打印译文，并与原始的前、后缀合并
|
||
|
||
try:
|
||
translated_text = self.optimize_with_ollama(translated_text, ollama_url, ollama_model)
|
||
match = re.search(r'#(.*?)#', translated_text)
|
||
if match:
|
||
translated_text = match.group(1)
|
||
translated_text = re.sub(r'^[\s|,]+|[\s|,]+$', '', translated_text.strip())
|
||
|
||
prompt_text = ", ".join(filter(None, [prefix, translated_text, suffix]))
|
||
print(f" \033[94m{prefix}\033[0m, \033[90;1m{translated_text}\033[0m, \033[94m{suffix}\033[0m") # 打印 Ollama 对译文的处理结果，并与原始的前、后缀合并
|
||
except Exception as e:
|
||
print(f"Failed to optimize prompt with Ollama in translated text: {e}")
|
||
|
||
else:
|
||
translated = model.generate(**tokenizer(translate_part, return_tensors="pt", padding=True))
|
||
translated_text = re.sub(r'^[\s,.]+|[\s,.]+$', '', tokenizer.decode(translated[0], skip_special_tokens=True))
|
||
|
||
prompt_text = ", ".join(filter(None, [prefix, translated_text, suffix]))
|
||
print(f" \033[94m{prefix}\033[0m, \033[92m{translated_text}\033[0m, \033[94m{suffix}\033[0m") # 打印译文，保留原始前、后缀不变，并完成合并
|
||
|
||
if ollama_for_all_3_parts:
|
||
prompt_text = prompt_text.replace('!', '') # 去掉 `!`
|
||
try:
|
||
prompt_text = self.optimize_with_ollama(prompt_text, ollama_url, ollama_model)
|
||
match = re.search(r'#(.*?)#', prompt_text)
|
||
if match:
|
||
prompt_text = match.group(1)
|
||
prompt_text = re.sub(r'^[\s|,]+|[\s|,]+$', '', prompt_text.strip())
|
||
|
||
print(f" \033[90;1m{prompt_text}\033[0m") # 打印 Ollama 对合并后的提示词（包括前、后缀和译文）整体的处理结果
|
||
except Exception as e:
|
||
print(f"Failed to optimize prompt with Ollama in all 3 parts: {e}")
|
||
# 去掉前缀 prefix 为空时，开头会出现的`,`
|
||
prompt_text = prompt_text.lstrip(',')
|
||
|
||
else:
|
||
prompt_text = text.strip()
|
||
print(f" \033[94m{prompt_text}\033[0m") # 打印无需翻译的原始提示词
|
||
|
||
if ollama_for_part_of_original:
|
||
try:
|
||
# 使用正则提取 `!!` 之间的部分，并移除 `!!`
|
||
match = re.search(r'^(.*?)!(.*?)!(.*?)$', text)
|
||
if match:
|
||
prefix, ollama_part, suffix = [re.sub(r'^[\s|,]+|[\s|,]+$', '', part.strip()) for part in match.groups()]
|
||
# prefix, ollama_part, suffix = match.groups()
|
||
ollama_part = self.optimize_with_ollama(ollama_part.strip(), ollama_url, ollama_model)
|
||
ollama_match = re.search(r'#(.*?)#', ollama_part)
|
||
if ollama_match:
|
||
ollama_part = ollama_match.group(1)
|
||
ollama_part = re.sub(r'^[\s|,]+|[\s|,]+$', '', ollama_part.strip())
|
||
|
||
prompt_text = f"{prefix.strip()}, {ollama_part.strip()}, {suffix.strip()}"
|
||
print(f" \033[32m{prefix}\033[0m, \033[90;1m{ollama_part}\033[0m, \033[32m{suffix}\033[0m") # 打印 Ollama 对无需翻译的原始提示词的处理结果，并保留原始前、后缀
|
||
except Exception as e:
|
||
print(f"Failed to optimize prompt with Ollama in part of original prompt: {e}")
|
||
|
||
elif ollama_for_original:
|
||
prompt_text = text.replace('!', '') # 去掉 `!`
|
||
try:
|
||
prompt_text = self.optimize_with_ollama(prompt_text, ollama_url, ollama_model)
|
||
match = re.search(r'#(.*?)#', prompt_text)
|
||
if match:
|
||
prompt_text = match.group(1)
|
||
prompt_text = re.sub(r'^[\s|,]+|[\s|,]+$', '', prompt_text.strip())
|
||
|
||
print(f" \033[90;1m{prompt_text}\033[0m") # 打印无需翻译但需要 Ollama 处理的提示词
|
||
except Exception as e:
|
||
print(f"Failed to optimize prompt with Ollama in original prompt: {e}")
|
||
|
||
# 去掉前缀 prefix 为空时，开头会出现的`,`
|
||
prompt_text = prompt_text.lstrip(',')
|
||
|
||
# 最终输出
|
||
print(f" \033[32m{prompt_text}\033[0m") # 打印进入 CLIP 编码流程的提示词
|
||
|
||
tokens = clip.tokenize(prompt_text)
|
||
cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
|
||
|
||
return ([[cond, {"pooled_output": pooled}]], prompt_text)
|
||
|
||
RETURN_TYPES = ("CONDITIONING", "STRING",)
|
||
FUNCTION = "mtencode"
|
||
CATEGORY = "MTCLIPEncode"
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"MTCLIPEncode": MTCLIPEncode,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"MTCLIPEncode": "MTCLIPEncode",
|
||
} |