initial program, throws error without civitai info
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user