dogmantc / lora-wildcard-maker Goto Github PK

import os
import json

# Define directories 
directory = os.path.join(os.getcwd(),"models")
output_dir = os.path.join(os.getcwd(), "extensions", "sd-dynamic-prompts", "wildcards")

# Get words and loras from .civitai.info files
words = {}
subdirectories = set()  # Store the subdirectories

for root, dirs, files in os.walk(directory):
    origin = os.path.basename(root)  # Get the subdirectory name as the origin
    subdirectories.add(origin[0:200])  # Add the subdirectory to the set
    for filename in files:
        if filename.endswith(".civitai.info"):
            file_path = os.path.join(root, filename)
            with open(file_path, "r") as file:
                twords = "" #Base value
                file_contents = file.read()
                json_data = json.loads(file_contents)
                LoraPrompt = f"<lora:{filename.replace('.civitai.info','')}"
                trained_words = json_data.get("trainedWords", [])
                if trained_words:
                    twords = ",".join([str(item.strip().strip("\"'[{(").strip("\"' ]})")) for item in trained_words]) #add trained words
                if "images" in json_data:
                    Imagelist = json_data.get("images",[])
                    if Imagelist:
                        Meta = json_data['images'][0]['meta']
                        if Meta is not None:
                            if "prompt" in Meta:
                                prompt = json_data['images'][0]['meta']['prompt'] #get first image test prompt
                                prompt = prompt.replace('BREAK','')
                                prompt = prompt.replace('[','{')
                                prompt = prompt.replace(']','}')
                                prompt = prompt.replace('\n','')
                                prompt = prompt.replace('','')
                                prompt = prompt.replace('lyco','lora')
                                NumLora = prompt.find(LoraPrompt)
                                if NumLora < 1:
                                    LoraPrompt = f"<lora:{filename.replace('.civitai.info','')}:0.8>"
                                else:
                                    LoraPrompt = ""
                                words[file_path] = f"<REM:{filename.replace('.civitai.info','')}>, {LoraPrompt}, {prompt}%{origin}"  # Append origin to the lora
                            else:
                                LoraPrompt = f"<lora:{filename.replace('.civitai.info','')}:0.8>, {twords}"
                                words[file_path] = f"<REM:{filename.replace('.civitai.info','')}>, {LoraPrompt}%{origin}"  # Append origin to the lora
                        else:
                            LoraPrompt = f"<lora:{filename.replace('.civitai.info','')}:0.8>, {twords}"
                            words[file_path] = f"<REM:{filename.replace('.civitai.info','')}>, {LoraPrompt}%{origin}"  # Append origin to the lora
                    else:
                        LoraPrompt = f"<lora:{filename.replace('.civitai.info','')}:0.8>, {twords}"
                        words[file_path] = f"<REM:{filename.replace('.civitai.info','')}>, {LoraPrompt}%{origin}"  # Append origin to the lora
                else:
                    LoraPrompt = f"<lora:{filename.replace('.civitai.info','')}:0.8>, {twords}"
                    words[file_path] = f"<REM:{filename.replace('.civitai.info','')}>, {LoraPrompt}%{origin}"  # Append origin to the lora
# Write Overall txt file
loras_file_path = os.path.join(output_dir, "AllPrompts.txt")
with open(loras_file_path, "w", encoding="utf-8") as file:
    for word, thing in words.items():
        prompt_without_origin = thing.rsplit('%', 1)[0]  # Remove the "-{origin}" from the line
        file.write(f"{prompt_without_origin}\n")
# Write subdir Prompt_{subdirectory}.txt files and remove "-{origin}" from lines
for subdir in subdirectories:
    if "lora" in subdir:
        sub_output_file = os.path.join(output_dir, f"Prompts_{subdir}.txt")
        with open(sub_output_file, "w", encoding="utf-8") as file:
            for word, thing in words.items():
                if thing.endswith(f"%{subdir}"):  # Check if the line belongs to the current subdirectory
                    prompt_without_origin = thing.rsplit('%', 1)[0]  # Remove the "-{origin}" from the line
                    file.write(f"{prompt_without_origin}\n")
    if "hypernetwork" in subdir:
        sub_output_file = os.path.join(output_dir, f"Prompts_{subdir}.txt")
        with open(sub_output_file, "w", encoding="utf-8") as file:
            for word, thing in words.items():
                if thing.endswith(f"%{subdir}"):  # Check if the line belongs to the current subdirectory
                    prompt_without_origin = thing.rsplit('%', 1)[0]  # Remove the "-{origin}" from the line
                    file.write(f"{prompt_without_origin}\n")
    if "Stable-diffusion" in subdir:
        sub_output_file = os.path.join(output_dir, f"Prompts_{subdir}.txt")
        with open(sub_output_file, "w", encoding="utf-8") as file:
            for word, thing in words.items():
                if thing.endswith(f"%{subdir}"):  # Check if the line belongs to the current subdirectory
                    prompt_without_origin = thing.rsplit('%', 1)[0]  # Remove the "-{origin}" from the line
                    file.write(f"{prompt_without_origin}\n")