main5.py

import os
import zipfile
import ipywidgets as widgets
from IPython.display import display
from tqdm import tqdm

class ModelExtractor:
    def __init__(self, base_path):
        self.base_path = base_path
        self.vae_zip_file_path = os.path.join(base_path, "models/Stable-diffusion/\u805A\u661Fvae.zip")
        self.vae_extract_path = os.path.join(base_path, "models/VAE/\u805A\u661FVAE")
        self.embeddings_zip_file_path = os.path.join(base_path, "models/Lora/\u805A\u661Fembeddings.zip")
        self.embeddings_extract_path = os.path.join(base_path, "embeddings/\u805A\u661Fembeddings")
        self.lora_zip_file_path = os.path.join(base_path, "models/Stable-diffusion/\u805A\u661Flora.zip")
        self.lora_extract_path = os.path.join(base_path, "models/Lora/\u805A\u661FLora")

        # 设置文件系统编码环境
        os.environ['LANG'] = 'en_US.UTF-8'

        # 创建输入框和按钮
        self.old_char_input = widgets.Text(value='/mnt/8BxdTF', description='替换前的字符:')
        self.new_char_input = widgets.Text(value='', description='替换后的字符:')
        self.replace_button = widgets.Button(description='替换')

        # 显示输入框和按钮
        display(self.old_char_input, self.new_char_input, self.replace_button)

        # 将按钮与回调函数连接
        self.replace_button.on_click(self.replace_callback)

    def find_and_replace_path(self, file_path, old_str, new_str):
        return file_path.replace(old_str, new_str)

    def unzip_and_fix_encoding(self, zip_file, extract_path, old_char, new_char):
        # 解压文件
        with zipfile.ZipFile(zip_file, 'r') as zip_ref:
            zip_ref.extractall(extract_path)

        # 删除原始压缩文件
        os.remove(zip_file)

    def replace_callback(self, button):
        # 获取输入框中的值
        old_char = self.old_char_input.value
        new_char = self.new_char_input.value

        # 替换路径
        self.base_path = self.find_and_replace_path(self.base_path, old_char, new_char)
        self.vae_zip_file_path = self.find_and_replace_path(self.vae_zip_file_path, old_char, new_char)
        self.vae_extract_path = self.find_and_replace_path(self.vae_extract_path, old_char, new_char)

        self.embeddings_zip_file_path = self.find_and_replace_path(self.embeddings_zip_file_path, old_char, new_char)
        self.embeddings_extract_path = self.find_and_replace_path(self.embeddings_extract_path, old_char, new_char)

        self.lora_zip_file_path = self.find_and_replace_path(self.lora_zip_file_path, old_char, new_char)
        self.lora_extract_path = self.find_and_replace_path(self.lora_extract_path, old_char, new_char)

        # 替换文件夹路径
        self.vae_extract_path = self.find_and_replace_path(self.vae_extract_path, old_char, new_char)
        self.embeddings_extract_path = self.find_and_replace_path(self.embeddings_extract_path, old_char, new_char)
        self.lora_extract_path = self.find_and_replace_path(self.lora_extract_path, old_char, new_char)

        # 解压 VAE 并修复文件名编码
        self.unzip_and_fix_encoding(self.vae_zip_file_path, self.vae_extract_path, old_char, new_char)

        # 解压 embeddings 并修复文件名编码
        self.unzip_and_fix_encoding(self.embeddings_zip_file_path, self.embeddings_extract_path, old_char, new_char)

        # 解压 Lora 并修复文件名编码
        self.unzip_and_fix_encoding(self.lora_zip_file_path, self.lora_extract_path, old_char, new_char)

        # 遍历并修复两个路径的文件名编码
        for current_base_path in [self.vae_extract_path, self.lora_extract_path]:
            for root, dirs, files in os.walk(current_base_path):
                for file in files:
                    # 拼接文件的完整路径
                    file_path = os.path.join(root, file)
                    # 从乱码中解析文件名
                    decoded_name = file.encode('cp437').decode('gbk')
                    # 更改文件名
                    os.rename(file_path, os.path.join(root, decoded_name))

            print(f"{current_base_path}路径的文件名编码修复完成。")

        print("VAE、embeddings 和 Lora 解压完成，并进行字符替换和文件名编码修复，原始压缩文件已删除.")

# 实例化类并执行替换操作
model_extractor = ModelExtractor("/mnt/8BxdTF")