Add support to other img format, fixed dropbox update

This commit is contained in:
MalumaDev 2022-10-16 00:01:00 +02:00
parent 3f5c3b981e
commit 3d21684ee3

View File

@ -8,7 +8,7 @@ import gradio as gr
import torch import torch
from PIL import Image from PIL import Image
from modules import shared from modules import shared
from modules.shared import device, aesthetic_embeddings from modules.shared import device
from transformers import CLIPModel, CLIPProcessor from transformers import CLIPModel, CLIPProcessor
from tqdm.auto import tqdm from tqdm.auto import tqdm
@ -20,7 +20,7 @@ def get_all_images_in_folder(folder):
def check_is_valid_image_file(filename): def check_is_valid_image_file(filename):
return filename.lower().endswith(('.png', '.jpg', '.jpeg')) return filename.lower().endswith(('.png', '.jpg', '.jpeg', ".gif", ".tiff", ".webp"))
def batched(dataset, total, n=1): def batched(dataset, total, n=1):
@ -73,6 +73,6 @@ def generate_imgs_embd(name, folder, batch_size):
Aesthetic embedding saved to {html.escape(path)} Aesthetic embedding saved to {html.escape(path)}
""" """
shared.update_aesthetic_embeddings() shared.update_aesthetic_embeddings()
return gr.Dropdown(sorted(aesthetic_embeddings.keys()), label="Imgs embedding", return gr.Dropdown.update(choices=sorted(shared.aesthetic_embeddings.keys()), label="Imgs embedding",
value=sorted(aesthetic_embeddings.keys())[0] if len( value=sorted(shared.aesthetic_embeddings.keys())[0] if len(
aesthetic_embeddings) > 0 else None), res, "" shared.aesthetic_embeddings) > 0 else None), res, ""