MergeTools.py

import os
import shutil
from datetime import datetime
import pytz
import subprocess
import argparse

# 设置时区
beijing_tz = pytz.timezone('Asia/Shanghai')
# 获取当前北京时间
beijing_time = datetime.now(beijing_tz)
# 当前日期字符串
date_str = beijing_time.strftime('%m%d')
# date_str =('0818')
# 目标基目录
target_base_dir = f'/mnt/workspace/{date_str}/'

def create_directories():
    # 创建目标基目录
    os.makedirs(target_base_dir, exist_ok=True)
    # 生成十个子文件夹
    for i in range(1, 11):
        subdir = f'{i:02d}'
        target_dir = os.path.join(target_base_dir, subdir)
        # 创建子文件夹（如果已经存在则跳过）
        if not os.path.exists(target_dir):
            os.makedirs(target_dir)
            print(f'创建文件夹: {target_dir}')
        else:
            print(f'文件夹已存在: {target_dir}')
        
        # 检查并将 /mnt/workspace/description.txt 复制到 target_dir
        description_src = '/mnt/workspace/description.txt'
        description_dest = os.path.join(target_base_dir, 'description.txt')
        if os.path.exists(description_src):
            shutil.copy(description_src, description_dest)
            print(f'复制文件: {description_src} 到 {description_dest}')
        else:
            print(f'文件不存在: {description_src}')

def move_files_by_date():
    # 第一步先检查 target_base_dir 是否存在，如果不存在则 print("WRONG")
    if not os.path.exists(target_base_dir):
        print("WRONG")
        return
    
    source_dir = '/mnt/workspace/stable-diffusion-webui/models/Lora/'
    
    # 循环处理序号 01 到 10
    for i in range(1, 11):
        seq_num = f'{i:02d}'  # 将序号格式化为两位数
        source_file = os.path.join(source_dir, f'{date_str}_{seq_num}.safetensors')
        target_dir = os.path.join(target_base_dir, seq_num)
        
        # 如果源文件存在，则移动到目标目录
        if os.path.exists(source_file):
            os.makedirs(target_dir, exist_ok=True)
            shutil.move(source_file, target_dir)
            print(f'{source_file} 移动到 {target_dir}')
        else:
            print(f'{source_file} 不存在')

def packup():
    zip_command = ['7z', 'a', f'/mnt/workspace/LoraUpload/{date_str}.zip', f'/mnt/workspace/{date_str}']
    subprocess.run(zip_command)

def move_upscale():
    # 第一步先检查 target_base_dir 是否存在，如不存在则 create_directories()
    if not os.path.exists(target_base_dir):
        create_directories()
    
    target_to_move = None
    
    # 第二步依次按顺序检查 target_base_dir 下的 01 到 10 这十个文件夹
    for i in range(1, 11):
        subdir = f'{i:02d}'
        check_dir = os.path.join(target_base_dir, subdir)
        upscale_dir = os.path.join(check_dir, 'Like_Upscale')
        if not os.path.exists(upscale_dir):
            target_to_move = check_dir
            break
    
    # 如果所有文件夹都包含 Upscale 文件夹，target_to_move 为 FULL
    if target_to_move is None:
        print("FULL")
    else:
        # 第三步，将 /mnt/workspace/Photo/Like_Upscale 移动到 target_to_move 文件夹里
        source_upscale = '/mnt/workspace/Photo/Like_Upscale'
        if os.path.exists(source_upscale):
            shutil.move(source_upscale, target_to_move)
            print(f'{source_upscale} 移动到 {target_to_move}')
        else:
            print(f'{source_upscale} 不存在')

def clear_LoraUpload_model():
    lora_upload_dir = '/mnt/workspace/LoraUpload'
    # 获取当前日期
    today_str = datetime.now(beijing_tz).strftime('%m%d')

    if os.path.exists(lora_upload_dir):
        # 列出所有.zip文件，并按日期排序
        zip_files = sorted(
            [f for f in os.listdir(lora_upload_dir) if f.endswith('.zip')],
            reverse=True
        )
        
        # 保留最新的一个.zip文件
        for zip_file in zip_files[1:]:  # 跳过最新的
            zip_path = os.path.join(lora_upload_dir, zip_file)
            os.remove(zip_path)
            print(f'已删除: {zip_path}')
    else:
        print(f'{lora_upload_dir} 不存在')

def upload_to_modelscope():
    from modelscope.hub.api import HubApi
    YOUR_ACCESS_TOKEN = '4654205d-1369-44e3-b3ca-d12678e2950d'
    api = HubApi()
    api.login(YOUR_ACCESS_TOKEN)
    api.push_model(
        model_id="LIBLIBAI/LoraMerge", 
        model_dir="/mnt/workspace/LoraUpload"
    )
    
def ClearLike():
    path = "/mnt/workspace/Photo/Like"
    if os.path.exists(path):
        shutil.rmtree(path)
        print(f"{path} 已删除")
    else:
        print(f"{path} 不存在")
        
def move_and_delete_like_upscale():
    base_dir = target_base_dir
    
    # 遍历 01 到 10 子文件夹
    for i in range(1, 11):
        subdir = f'{i:02d}'
        like_upscale_dir = os.path.join(base_dir, subdir, 'Like_Upscale')
        
        # 检查 Like_Upscale 文件夹是否存在
        if os.path.exists(like_upscale_dir) and os.path.isdir(like_upscale_dir):
            # 遍历 Like_Upscale 文件夹中的所有文件
            for file_name in os.listdir(like_upscale_dir):
                file_path = os.path.join(like_upscale_dir, file_name)
                
                # 检查是否为文件（而非子目录）
                if os.path.isfile(file_path):
                    # 移动文件到父级目录
                    shutil.move(file_path, os.path.join(base_dir, subdir))
                    print(f'已移动: {file_path} 到 {os.path.join(base_dir, subdir)}')
            
            # 删除 Like_Upscale 文件夹
            os.rmdir(like_upscale_dir)
            print(f'已删除文件夹: {like_upscale_dir}')
        else:
            print(f'{like_upscale_dir} 不存在或不是文件夹')

# 主程序入口
if __name__ == "__main__":
    # 创建 ArgumentParser 对象
    parser = argparse.ArgumentParser(description='执行特定函数的脚本')
    
    # 添加命令行参数
    parser.add_argument('--create_directories', action='store_true', help='执行 create_directories 函数')
    parser.add_argument('--move_and_delete_like_upscale', action='store_true', help='执行 move_and_delete_like_upscale 函数')
    parser.add_argument('--ClearLike', action='store_true', help='执行 ClearLike 函数')
    parser.add_argument('--move_files_by_date', action='store_true', help='执行 move_files_by_date 函数')
    parser.add_argument('--packup', action='store_true', help='执行 packup 函数')
    parser.add_argument('--move_upscale', action='store_true', help='执行 move_upscale 函数')
    parser.add_argument('--clear_LoraUpload_model', action='store_true', help='执行 clear_LoraUpload_model 函数')
    parser.add_argument('--upload_to_modelscope', action='store_true', help='执行 upload_to_modelscope 函数')

    # 解析参数
    args = parser.parse_args()

    # 根据传入的参数执行对应的函数
    if args.create_directories:
        create_directories()
    if args.ClearLike:
        ClearLike()
    if args.move_files_by_date:
        move_files_by_date()
    if args.packup:
        packup()
    if args.move_upscale:
        move_upscale()
    if args.move_and_delete_like_upscale:
        move_and_delete_like_upscale()
    if args.clear_LoraUpload_model:
        clear_LoraUpload_model()
    if args.upload_to_modelscope:
        upload_to_modelscope()