244 lines
9.5 KiB
Python
244 lines
9.5 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""
|
|||
|
|
generate_assets_local.py — 本地ComfyUI批量生成角色/场景/道具资产图
|
|||
|
|
替代线上Seedream API（¥0.15-1/张）→ 本地ComfyUI（¥0电费）
|
|||
|
|
用法:
|
|||
|
|
python tools/generate_assets_local.py <项目目录> [--comfy-url URL]
|
|||
|
|
python tools/generate_assets_local.py <项目目录> --list # 列出待生成的资产
|
|||
|
|
"""
|
|||
|
|
import sys, os, json, argparse, subprocess
|
|||
|
|
|
|||
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|||
|
|
|
|||
|
|
# 默认资产生成prompt模板
|
|||
|
|
ASSET_PROMPTS = {
|
|||
|
|
"CHAR": {
|
|||
|
|
"template": "{name}, {description}, {style}, full body character design, character sheet, consistent clothing, clean background, high quality, detailed",
|
|||
|
|
"default_style": "digital art, anime style, vibrant colors"
|
|||
|
|
},
|
|||
|
|
"ENV": {
|
|||
|
|
"template": "{name}, {description}, {style}, wide angle view, environmental design, detailed scene, high quality",
|
|||
|
|
"default_style": "digital painting, cinematic lighting, detailed environment"
|
|||
|
|
},
|
|||
|
|
"PROP": {
|
|||
|
|
"template": "{name}, {description}, {style}, isolated object, clean background, detailed, high quality",
|
|||
|
|
"default_style": "product photography, sharp focus, detailed texture"
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def generate_workflow(prompt, asset_type, asset_name, output_prefix="asset"):
|
|||
|
|
"""生成ComfyUI API调用JSON"""
|
|||
|
|
workflow = {
|
|||
|
|
"3": {
|
|||
|
|
"class_type": "KSampler",
|
|||
|
|
"inputs": {
|
|||
|
|
"seed": abs(hash(asset_name)) % 1000000,
|
|||
|
|
"steps": 25,
|
|||
|
|
"cfg": 7,
|
|||
|
|
"sampler_name": "euler",
|
|||
|
|
"scheduler": "normal",
|
|||
|
|
"denoise": 1,
|
|||
|
|
"model": ["4", 0],
|
|||
|
|
"positive": ["6", 0],
|
|||
|
|
"negative": ["7", 0],
|
|||
|
|
"latent_image": ["5", 0]
|
|||
|
|
}
|
|||
|
|
},
|
|||
|
|
"4": {
|
|||
|
|
"class_type": "CheckpointLoaderSimple",
|
|||
|
|
"inputs": {"ckpt_name": "realvisxlV40_v40BDFPonNoobVae.safetensors"}
|
|||
|
|
},
|
|||
|
|
"5": {
|
|||
|
|
"class_type": "EmptyLatentImage",
|
|||
|
|
"inputs": {"width": 1024, "height": 1024, "batch_size": 1}
|
|||
|
|
},
|
|||
|
|
"6": {
|
|||
|
|
"class_type": "CLIPTextEncode",
|
|||
|
|
"inputs": {"text": prompt, "clip": ["4", 1]}
|
|||
|
|
},
|
|||
|
|
"7": {
|
|||
|
|
"class_type": "CLIPTextEncode",
|
|||
|
|
"inputs": {"text": "worst quality, low quality, blurry, deformed", "clip": ["4", 1]}
|
|||
|
|
},
|
|||
|
|
"8": {
|
|||
|
|
"class_type": "VAEDecode",
|
|||
|
|
"inputs": {"samples": ["3", 0], "vae": ["4", 2]}
|
|||
|
|
},
|
|||
|
|
"9": {
|
|||
|
|
"class_type": "SaveImage",
|
|||
|
|
"inputs": {"filename_prefix": f"{output_prefix}_{asset_name}", "images": ["8", 0]}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return workflow
|
|||
|
|
|
|||
|
|
def load_asset_list(project_dir):
|
|||
|
|
"""加载项目资产清单"""
|
|||
|
|
# 检查是否有 ASSET-LIST 文件
|
|||
|
|
for pattern in ['ASSET-LIST*', 'assets*.json', 'CHARACTERS*']:
|
|||
|
|
import glob
|
|||
|
|
matches = glob.glob(os.path.join(project_dir, pattern))
|
|||
|
|
if matches:
|
|||
|
|
with open(matches[0], 'r', encoding='utf-8') as f:
|
|||
|
|
return json.load(f)
|
|||
|
|
|
|||
|
|
# 检查项目protocols目录
|
|||
|
|
protocols_dir = os.path.join(project_dir, 'protocols')
|
|||
|
|
if os.path.isdir(protocols_dir):
|
|||
|
|
for fname in os.listdir(protocols_dir):
|
|||
|
|
if 'ASSET' in fname.upper() or 'CHAR' in fname.upper():
|
|||
|
|
with open(os.path.join(protocols_dir, fname), 'r', encoding='utf-8') as f:
|
|||
|
|
content = f.read()
|
|||
|
|
try:
|
|||
|
|
return json.loads(content)
|
|||
|
|
except:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def scan_project_assets(project_dir):
|
|||
|
|
"""扫描项目目录，识别需要生成的资产"""
|
|||
|
|
assets = []
|
|||
|
|
|
|||
|
|
# 检查 characters/ envs/ props/ 目录
|
|||
|
|
for asset_type, subdir in [('CHAR', 'characters'), ('ENV', 'envs'), ('PROP', 'props')]:
|
|||
|
|
asset_dir = os.path.join(project_dir, 'assets', subdir)
|
|||
|
|
if os.path.isdir(asset_dir):
|
|||
|
|
for f in os.listdir(asset_dir):
|
|||
|
|
if f.endswith(('.json', '.hdlp', '.txt')):
|
|||
|
|
with open(os.path.join(asset_dir, f), 'r', encoding='utf-8') as fh:
|
|||
|
|
content = fh.read()
|
|||
|
|
assets.append({
|
|||
|
|
'type': asset_type,
|
|||
|
|
'name': os.path.splitext(f)[0],
|
|||
|
|
'description': content[:200],
|
|||
|
|
'file': os.path.join(asset_dir, f)
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 检查分镜JSON中出现的角色/场景
|
|||
|
|
for root, dirs, files in os.walk(project_dir):
|
|||
|
|
for f in files:
|
|||
|
|
if f.startswith('STORYBOARD') and f.endswith('.json'):
|
|||
|
|
try:
|
|||
|
|
with open(os.path.join(root, f), 'r', encoding='utf-8') as fh:
|
|||
|
|
sb = json.load(fh)
|
|||
|
|
for shot in sb.get('shots', []):
|
|||
|
|
for char in shot.get('characters', []):
|
|||
|
|
if not any(a['name'] == char for a in assets):
|
|||
|
|
assets.append({
|
|||
|
|
'type': 'CHAR',
|
|||
|
|
'name': char,
|
|||
|
|
'description': f'角色 {char}',
|
|||
|
|
'source': f
|
|||
|
|
})
|
|||
|
|
for scene in shot.get('scenes', []):
|
|||
|
|
if not any(a['name'] == scene for a in assets):
|
|||
|
|
assets.append({
|
|||
|
|
'type': 'ENV',
|
|||
|
|
'name': scene,
|
|||
|
|
'description': f'场景 {scene}',
|
|||
|
|
'source': f
|
|||
|
|
})
|
|||
|
|
except:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
return assets
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser(description='本地ComfyUI批量生成资产图')
|
|||
|
|
parser.add_argument('project_dir', help='项目目录')
|
|||
|
|
parser.add_argument('--comfy-url', default='http://127.0.0.1:8188', help='ComfyUI API地址')
|
|||
|
|
parser.add_argument('--style', default='', help='全局风格')
|
|||
|
|
parser.add_argument('--list', action='store_true', help='只列出待生成资产')
|
|||
|
|
parser.add_argument('--output-dir', help='资产输出目录 (默认: project_dir/assets/)')
|
|||
|
|
parser.add_argument('--dry-run', action='store_true', help='只生成workflow JSON，不调API')
|
|||
|
|
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
if not os.path.isdir(args.project_dir):
|
|||
|
|
print(f'❌ 项目目录不存在: {args.project_dir}')
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
# 扫描资产
|
|||
|
|
assets = scan_project_assets(args.project_dir)
|
|||
|
|
|
|||
|
|
if not assets:
|
|||
|
|
print('⚠️ 未发现需要生成的资产')
|
|||
|
|
print(' 提示: 将角色描述放在 assets/characters/ 目录下')
|
|||
|
|
print(' 或: 先运行 storyboard_to_workflow.py 生成分镜后再扫描')
|
|||
|
|
sys.exit(0)
|
|||
|
|
|
|||
|
|
# 去重
|
|||
|
|
seen = set()
|
|||
|
|
unique_assets = []
|
|||
|
|
for a in assets:
|
|||
|
|
key = f"{a['type']}-{a['name']}"
|
|||
|
|
if key not in seen:
|
|||
|
|
seen.add(key)
|
|||
|
|
unique_assets.append(a)
|
|||
|
|
assets = unique_assets
|
|||
|
|
|
|||
|
|
print(f'📦 发现 {len(assets)} 个待生成资产:')
|
|||
|
|
for a in assets:
|
|||
|
|
print(f' [{a["type"]}] {a["name"]:20s} {a.get("description", "")[:40]}')
|
|||
|
|
|
|||
|
|
if args.list:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 输出目录
|
|||
|
|
out_dir = args.output_dir or os.path.join(args.project_dir, 'assets')
|
|||
|
|
assets_char_dir = os.path.join(out_dir, 'characters')
|
|||
|
|
assets_env_dir = os.path.join(out_dir, 'envs')
|
|||
|
|
assets_prop_dir = os.path.join(out_dir, 'props')
|
|||
|
|
for d in [assets_char_dir, assets_env_dir, assets_prop_dir]:
|
|||
|
|
os.makedirs(d, exist_ok=True)
|
|||
|
|
|
|||
|
|
# 生成并调用
|
|||
|
|
for i, asset in enumerate(assets):
|
|||
|
|
atype = asset['type']
|
|||
|
|
name = asset['name']
|
|||
|
|
desc = asset.get('description', name)
|
|||
|
|
style = args.style or ASSET_PROMPTS.get(atype, {}).get('default_style', '')
|
|||
|
|
template = ASSET_PROMPTS.get(atype, {}).get('template', '{description}, {style}')
|
|||
|
|
|
|||
|
|
prompt = template.format(name=name, description=desc, style=style)
|
|||
|
|
|
|||
|
|
print(f'\n🎨 [{i+1}/{len(assets)}] {name} ({atype})')
|
|||
|
|
print(f' Prompt: {prompt[:80]}...')
|
|||
|
|
|
|||
|
|
workflow = generate_workflow(prompt, atype, name)
|
|||
|
|
|
|||
|
|
# 确定输出位置
|
|||
|
|
dir_map = {'CHAR': assets_char_dir, 'ENV': assets_env_dir, 'PROP': assets_prop_dir}
|
|||
|
|
asset_out_dir = dir_map.get(atype, out_dir)
|
|||
|
|
|
|||
|
|
wf_path = os.path.join(asset_out_dir, f'_{name}_workflow.json')
|
|||
|
|
with open(wf_path, 'w', encoding='utf-8') as f:
|
|||
|
|
json.dump(workflow, f, ensure_ascii=False, indent=2)
|
|||
|
|
|
|||
|
|
if args.dry_run:
|
|||
|
|
print(f' ✅ 工作流JSON: {wf_path}')
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 调ComfyUI API
|
|||
|
|
print(f' 📤 发送ComfyUI...')
|
|||
|
|
try:
|
|||
|
|
resp = subprocess.run([
|
|||
|
|
'curl', '-s', '-X', 'POST',
|
|||
|
|
f'{args.comfy_url}/prompt',
|
|||
|
|
'-H', 'Content-Type: application/json',
|
|||
|
|
'-d', json.dumps({"prompt": workflow})
|
|||
|
|
], capture_output=True, text=True, timeout=30)
|
|||
|
|
|
|||
|
|
result = json.loads(resp.stdout) if resp.stdout else {}
|
|||
|
|
if 'error' in result:
|
|||
|
|
print(f' ❌ API错误: {result["error"]}')
|
|||
|
|
else:
|
|||
|
|
print(f' ✅ 已提交 (task: {result.get("task_id", "?")})')
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f' ⚠️ 调用失败: {e}')
|
|||
|
|
|
|||
|
|
print(f'\n✅ 完成! 资产将保存到: {out_dir}')
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
main()
|