initial program, throws error without civitai info

This commit is contained in:
2023-06-22 09:20:58 -04:00
parent 4226f47c44
commit ddfe682499
3 changed files with 165 additions and 1 deletions

View File

@@ -0,0 +1,68 @@
import glob
import io
import json
import os
from os import path
from datetime import datetime
import PySimpleGUI as sg
from PIL import Image
def get_civitai_text(civitai_path: str) -> str:
with open(civitai_path, "r") as f:
civitai_json = json.load(f)
created_at: datetime = datetime.fromisoformat(civitai_json["createdAt"])
return f"""Model Name: {civitai_json["model"]["name"]}
Checkpoint Name: {civitai_json["name"]}
Created at: {created_at}"""
def get_extra_filenames(model_path: str):
base, ext = os.path.splitext(model_path)
preview_image: str = base + ".preview.png"
civitai_info: str = base + ".civitai.info"
return preview_image if path.exists(preview_image) else None, civitai_info if path.exists(civitai_info) else None
def get_models_from_folder(model_folder: str):
return glob.glob(path.join(model_folder, "**/*.safetensors"), recursive=True) + glob.glob(
path.join(model_folder, "**/*.ckpt"), recursive=True)
layout = [
[sg.Text("Enter a folder:")],
[sg.InputText(key="model_folder", enable_events=True), sg.FolderBrowse()],
[sg.Combo(key="main_folder", enable_events=True, values=[], size=60),
sg.Combo(key="models", enable_events=True, values=[], size=60)],
[sg.Image(key="preview_image", enable_events=True)],
[sg.Text(key="civitai_info", enable_events=True)],
]
window = sg.Window("Model Folder", layout)
while True:
event, values = window.read()
print(event)
if event == sg.WIN_CLOSED or event == "Cancel":
break
if event == "model_folder":
model_folder = values["model_folder"]
if model_folder:
list_subfolders_with_paths = [f.path for f in os.scandir(model_folder) if f.is_dir()]
if list_subfolders_with_paths:
window["main_folder"].update(values=list_subfolders_with_paths)
if event == "main_folder":
model_subfolder = values["main_folder"]
models = get_models_from_folder(model_subfolder)
if model_subfolder and models:
window["models"].update(values=models)
if event == "models":
model_path = values["models"]
preview_image, civitai_info = get_extra_filenames(model_path)
img = Image.open(preview_image)
img.thumbnail((512, 512))
bio = io.BytesIO()
img.save(bio, format="PNG")
window["preview_image"].update(data=bio.getvalue())
window["civitai_info"].update(get_civitai_text(civitai_info))
window.close()