lora merger

scripts/data_process.py

@@ -0,0 +1,85 @@
+import torch, os, dashscope
+import pandas as pd
+from tqdm import tqdm
+from diffsynth import load_state_dict, hash_state_dict_keys
+
+
+def search_for_model_file(path, allow_file_extensions=(".safetensors",)):
+    for file_name in os.listdir(path):
+        for file_extension in allow_file_extensions:
+            if file_name.endswith(file_extension):
+                return os.path.join(path, file_name)
+            
+            
+def search_for_cover_images(path, allow_file_extensions=(".png", ".jpg", ".jpeg")):
+    image_files = []
+    for file_name in os.listdir(path):
+        for file_extension in allow_file_extensions:
+            if file_name.endswith(file_extension):
+                image_files.append(os.path.join(path, file_name))
+                break
+    return image_files
+
+
+def search_for_lora_data(path):
+    model_file = search_for_model_file(path)
+    if "_cover_images_" not in os.listdir(path):
+        return None
+    image_files = search_for_cover_images(os.path.join(path, "_cover_images_"))
+    if model_file is None or len(image_files) == 0:
+        return None
+    state_dict = load_state_dict(model_file)
+    if hash_state_dict_keys(state_dict, with_shape=False) != "52544ae3076666228978b738fbb8b086":
+        return None
+    return model_file, image_files
+
+
+def image_to_text(images=[], prompt="", system_prompt=None):
+    dashscope.api_key = "xxxxx" # TODO
+    messages = []
+    if system_prompt is not None:
+        messages.append({"role": "system", "content": system_prompt})
+    if not isinstance(images, list):
+        images = [images]
+    messages.append({"role": "user", "content": [{"text": prompt}] + [{"image": image} for image in images]})
+    response = dashscope.MultiModalConversation.call(model="qwen-vl-max-latest", messages=messages)
+    response = response["output"]["choices"][0]["message"]["content"][0]["text"]
+    return response
+
+
+qwen_i2t_prompt = '''
+You are a professional image captioner. 
+Generate a caption according to the image so that another image generation model can generate the image via the caption. Just return the string description, do not return anything else.
+'''.strip()
+
+
+def data_to_csv(model_file_list, image_file_list, text_list, save_path):
+    data_df = pd.DataFrame()
+    data_df["model_file"] = model_file_list
+    data_df["image_file"] = image_file_list
+    data_df["text"] = text_list
+    data_df.to_csv(save_path, index=False, encoding="utf-8-sig")
+
+
+base_path = "/data/zhiwen/LoRA-Fusion/models/FLUXLoRA"
+
+model_file_list = []
+image_file_list = []
+text_list = []
+
+for lora_name in tqdm(os.listdir(base_path)):
+    lora_folder_path = os.path.join(base_path, lora_name)
+    if os.path.isdir(lora_folder_path):
+        data = search_for_lora_data(lora_folder_path)
+        if data is not None:
+            model_file, image_files = data
+            for image_file in image_files:
+                try:
+                    text = image_to_text(image_file, prompt=qwen_i2t_prompt)
+                except:
+                    continue
+                model_file_list.append(model_file)
+                image_file_list.append(image_file)
+                text_list.append(text)
+                data_to_csv(model_file_list, image_file_list, text_list, "data/loras.csv")
+