|
import glob |
|
import json |
|
import logging |
|
import os |
|
import re |
|
import shutil |
|
from pathlib import Path |
|
|
|
from animatediff import get_dir |
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
data_dir = get_dir("data") |
|
|
|
extra_loading_regex = r'(<[^>]+?>)' |
|
|
|
def generate_config_from_civitai_info( |
|
lora_dir:Path, |
|
config_org:Path, |
|
out_dir:Path, |
|
lora_weight:float, |
|
): |
|
lora_abs_dir = lora_dir.absolute() |
|
config_org = config_org.absolute() |
|
out_dir = out_dir.absolute() |
|
|
|
civitais = sorted(glob.glob( os.path.join(lora_abs_dir, "*.civitai.info"), recursive=False)) |
|
|
|
with open(config_org, "r") as cf: |
|
org_config = json.load(cf) |
|
|
|
for civ in civitais: |
|
|
|
logger.info(f"convert {civ}") |
|
|
|
with open(civ, "r") as f: |
|
|
|
name = os.path.splitext(os.path.splitext(os.path.basename(civ))[0])[0] |
|
|
|
output_path = out_dir.joinpath(name + ".json") |
|
|
|
if os.path.isfile(output_path): |
|
logger.info("already converted -> skip") |
|
continue |
|
|
|
if os.path.isfile( lora_abs_dir.joinpath(name + ".safetensors")): |
|
lora_path = os.path.relpath(lora_abs_dir.joinpath(name + ".safetensors"), data_dir) |
|
elif os.path.isfile( lora_abs_dir.joinpath(name + ".ckpt")): |
|
lora_path = os.path.relpath(lora_abs_dir.joinpath(name + ".ckpt"), data_dir) |
|
else: |
|
logger.info("lora file not found -> skip") |
|
continue |
|
|
|
info = json.load(f) |
|
|
|
if not info: |
|
logger.info(f"empty civitai info -> skip") |
|
continue |
|
|
|
if info["model"]["type"] not in ("LORA","lora"): |
|
logger.info(f"unsupported type {info['model']['type']} -> skip") |
|
continue |
|
|
|
new_config = org_config.copy() |
|
|
|
new_config["name"] = name |
|
|
|
new_prompt_map = {} |
|
new_n_prompt = "" |
|
new_seed = -1 |
|
|
|
|
|
raw_prompt_map = {} |
|
|
|
i = 0 |
|
for img_info in info["images"]: |
|
if img_info["meta"]: |
|
try: |
|
raw_prompt = img_info["meta"]["prompt"] |
|
except Exception as e: |
|
logger.info("missing prompt") |
|
continue |
|
|
|
raw_prompt_map[str(10000 + i*32)] = raw_prompt |
|
|
|
new_prompt_map[str(i*32)] = re.sub(extra_loading_regex, '', raw_prompt) |
|
|
|
if not new_n_prompt: |
|
try: |
|
new_n_prompt = img_info["meta"]["negativePrompt"] |
|
except Exception as e: |
|
new_n_prompt = "" |
|
if new_seed == -1: |
|
try: |
|
new_seed = img_info["meta"]["seed"] |
|
except Exception as e: |
|
new_seed = -1 |
|
|
|
i += 1 |
|
|
|
if not new_prompt_map: |
|
new_prompt_map[str(0)] = "" |
|
|
|
for k in raw_prompt_map: |
|
|
|
new_prompt_map[k] = raw_prompt_map[k] |
|
|
|
new_config["prompt_map"] = new_prompt_map |
|
new_config["n_prompt"] = [new_n_prompt] |
|
new_config["seed"] = [new_seed] |
|
|
|
new_config["lora_map"] = {lora_path.replace(os.sep,'/'):lora_weight} |
|
|
|
with open( out_dir.joinpath(name + ".json"), 'w') as wf: |
|
json.dump(new_config, wf, indent=4) |
|
logger.info("converted!") |
|
|
|
preview = lora_abs_dir.joinpath(name + ".preview.png") |
|
if preview.is_file(): |
|
shutil.copy(preview, out_dir.joinpath(name + ".preview.png")) |
|
|
|
|
|
|
