import os
from glob import glob

os.chdir('/media/external/tmp')

def width(s):
    if len(s)<82:
        s=s+" \x1b[33m"+":"*(82-len(s))
    return s+"\x1b[34m"

def fidth(s):
    if len(s)<87:
        s=s+" \x1b[33m"+":"*(87-len(s))
    return s+"\x1b[34m"

FILENAMES=glob('*')
FILESTATS=[ os.stat(x) for x in FILENAMES ]
FILEDATAS=[]
MSG="\x1b[33m"+":"*73
print(MSG)
for i,_ in enumerate(FILENAMES):
    try:
        if not os.path.isdir(_):
            if FILESTATS[i].st_size >= 512:
                f=open(_,'rb')
                d=f.read(512)
                f.close()
                FILEDATAS.append(d)
                if d[3:7]==b"\x00\x00\x00\x00":
                    if d[8:91]==b'{"__metadata__":{"format":"pt"},"cond_stage_model.transformer.text_model.embeddings':
                        if FILESTATS[i].st_size > 200000000:
                            MSG="PASS\x1b[33m: \x1b[35mFILE IS A STABLE DIFFUSION PICKLED TENSOR CHECKPOINT MERGE"
                            print(f'\x1b[32m{width(MSG)} {_} - {FILESTATS[i].st_size}')
                        else:
                            GSM="\x1b[31mFAIL\x1b[33m: \x1b[35mFILE IS UHHH:"
                            print(f'\x1b[34m{fidth(GSM)} {_} - STAT: {FILESTATS[i].st_size} - {d[0:127]}')
                    elif d[8:50]==b'{"__metadata__":{"ss_batch_size_per_device':
                        if FILESTATS[i].st_size < 200000000:    
                            MSG="PASS\x1b[33m: \x1b[35mFILE IS A STABLE DIFFUSION PICKLED TENSOR LORA CHECKPOINT MERGE"
                            print(f'\x1b[32m{width(MSG)} {_} - {FILESTATS[i].st_size}')
                    elif d[8:55]==b'{"__metadata__":{"format":"pt","sd_merge_recipe':
                            MSG="PASS\x1b[33m: \x1b[35mFILE IS A STABLE DIFFUSION PICKLED TENSOR LORA CHECKPOINT MERGE"
                            print(f'\x1b[32m{width(MSG)} {_} - {FILESTATS[i].st_size}')
                    elif d[8:42]==b'{"__metadata__":{"ss_cache_latents':
                        if FILESTATS[i].st_size < 200000000:                    
                            MSG="PASS\x1b[33m: \x1b[35mFILE IS A STABLE DIFFUSION PICKLED TENSOR LORA"
                            print(f'\x1b[32m{width(MSG)} {_} - {FILESTATS[i].st_size}')
                    elif d[8:77]==b'{"__metadata__":{"format":"pt"},"cond_stage_model.model.ln_final.bias':
                        if FILESTATS[i].st_size > 200000000:                    
                            MSG="PASS\x1b[33m: \x1b[35mFILE IS A STABLE DIFFUSION PICKLED TENSOR CHECKPOINT TRAINED"
                            print(f'\x1b[32m{width(MSG)} {_} - {FILESTATS[i].st_size}')
                    elif d[8:47]==b'{"alphas_cumprod":{"dtype":"F32","shape':
                        if FILESTATS[i].st_size > 200000000:                    
                            MSG="PASS\x1b[33m: \x1b[35mFILE IS A STABLE DIFFUSION PICKLED TENSOR CHECKPOINT MERGE"
                            print(f'\x1b[32m{width(MSG)} {_} - {FILESTATS[i].st_size}')
                    elif d[8:73]==b'{"cond_stage_model.transformer.text_model.embeddings.position_ids':
                        if FILESTATS[i].st_size > 200000000:                    
                            MSG="PASS\x1b[33m: \x1b[35mFILE IS A STABLE DIFFUSION PICKLED TENSOR CHECKPOINT MERGE"
                            print(f'\x1b[32m{width(MSG)} {_} - {FILESTATS[i].st_size}')
                    elif d[0:63]==b'\x0bA\x01\x00\x00\x00\x00\x00{"__metadata__":{"ss_network_alpha":"4","ss_network_dim':
                        if FILESTATS[i].st_size > 200000000:                    
                            MSG="PASS\x1b[33m: \x1b[35mFILE IS A STABLE DIFFUSION CHECKPOINT MERGE"
                            print(f'\x1b[32m{width(MSG)} {_} - {FILESTATS[i].st_size}')
                    elif d[0:42]==b'\x9cD\x02\x00\x00\x00\x00\x00{"__metadata__":{"ss_network_alpha':
                        if FILESTATS[i].st_size > 200000000:                    
                            MSG="PASS\x1b[33m: \x1b[35mFILE IS A STABLE DIFFUSION LORA"
                            print(f'\x1b[32m{width(MSG)} {_} - {FILESTATS[i].st_size}')
                    elif d[0:59]==b'\xf8Q\x02\x00\x00\x00\x00\x00{"lora_te_text_model_encoder_layers_0_mlp_fc1.alpha':
                        if FILESTATS[i].st_size > 200000000:                    
                            MSG="PASS\x1b[33m: \x1b[35mFILE IS A STABLE DIFFUSION LORA"
                            print(f'\x1b[32m{width(MSG)} {_} - {FILESTATS[i].st_size}')
                    elif d[0:86]==b'\x1f\\\x02\x00\x00\x00\x00\x00{"cond_stage_model.transformer.text_model.embeddings.position_embedding.weight':
                        if FILESTATS[i].st_size > 200000000:                    
                            MSG="PASS\x1b[33m: \x1b[35mFILE IS A STABLE DIFFUSION PICKLED TENSOR CHECKPOINT LORA TRAINED"
                            print(f'\x1b[32m{width(MSG)} {_} - {FILESTATS[i].st_size}')
                    elif d[0:78]==b'\x00\x8b\x02\x00\x00\x00\x00\x00{"__metadata__":{"format":"pt"},"alphas_cumprod":{"dtype":"F32","shape':
                        if FILESTATS[i].st_size > 200000000:                    
                            MSG="PASS\x1b[33m: \x1b[35mFILE IS A STABLE DIFFUSION PICKLED TENSOR CHECKPOINT MERGE"
                            print(f'\x1b[32m{width(MSG)} {_} - {FILESTATS[i].st_size}')
                    else:
                        MSG="\x1b[31mFAIL\x1b[33m: \x1b[35mFILE IS NOT A CHECKPOINT OR LORA EMBEDDING"
                        print(f'\x1b[34m{fidth(MSG)} {_} - STAT: {FILESTATS[i].st_size} - {d[0:127]}')
                else:
                    if d[0:35]==b'PK\x03\x04\x00\x00\x08\x08\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0e\x00\x14\x00model':
                        if FILESTATS[i].st_size > 200000000:                    
                            MSG="PASS\x1b[33m: \x1b[35mFILE IS A STABLE DIFFUSION CHECKPOINT MERGE"
                            print(f'\x1b[32m{width(MSG)} {_} - {FILESTATS[i].st_size}')
                    elif d[0:64]==b'PK\x03\x04\x00\x00\x08\x08\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x10\x00\x12\x00archive/data.pklFB\x0e\x00ZZZZZZZZZZZZZZ':
                        if FILESTATS[i].st_size < 200000000:                    
                            MSG="PASS\x1b[33m: \x1b[35mFILE IS A STABLE DIFFUSION TEXTUAL INVERSION"
                            print(f'\x1b[32m{width(MSG)} {_} - {FILESTATS[i].st_size}')
                    else:                    
                        GSM="\x1b[31mFAIL\x1b[33m: \x1b[31mFILE IS NOT A PT FILE SIGNATURE"
                        print(f'\x1b[33m{fidth(GSM)} {_} - STAT {FILESTATS[i].st_size} - {d[0:127]}')
                #print('file read: break')
            else:
                GSM="\x1b[31mFAIL\x1b[33m: \x1b[31mFILESIZE IS TOO SMALL"
                print(f'\x1b[31m{fidth(GSM)} {_} - STAT: {FILESTATS[i].st_size} - {d[0:127]}')
    except Exception as e:
        GSM="\x1b[31mFAIL\x1b[33m: \x1b[31mGENERAL FAILURE"
        print(f'\x1b[31m{fidth(GSM)} {e} - FILE: {_} - STAT: {FILESTATS[i]} - {d[0:7]}')
MSG="\x1b[33m"+":"*73
print(MSG)