File size: 4,027 Bytes
ba96580 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import os
import folder_paths
import numpy as np
import torch
from PIL import Image
# Compatible with Alibaba EAS for quick launch
eas_cache_dir = '/stable-diffusion-cache/models'
# The directory of the cogvideoxfun
script_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
def tensor2pil(image):
return Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8))
def numpy2pil(image):
return Image.fromarray(np.clip(255. * image, 0, 255).astype(np.uint8))
def to_pil(image):
if isinstance(image, Image.Image):
return image
if isinstance(image, torch.Tensor):
return tensor2pil(image)
if isinstance(image, np.ndarray):
return numpy2pil(image)
raise ValueError(f"Cannot convert {type(image)} to PIL.Image")
def search_model_in_possible_folders(possible_folders, model):
model_name = None
# Check if the model exists in any of the possible folders within folder_paths.models_dir
for folder in possible_folders:
candidate_path = os.path.join(folder_paths.models_dir, folder, model)
if os.path.exists(candidate_path):
model_name = candidate_path
break
# If model_name is still None, check eas_cache_dir for each possible folder
if model_name is None and os.path.exists(eas_cache_dir):
for folder in possible_folders:
candidate_path = os.path.join(eas_cache_dir, folder, model)
if os.path.exists(candidate_path):
model_name = candidate_path
break
# If model_name is still None, prompt the user to download the model
if model_name is None:
print(f"Please download cogvideoxfun model to one of the following directories:")
for folder in possible_folders:
print(f"- {os.path.join(folder_paths.models_dir, folder)}")
if os.path.exists(eas_cache_dir):
print(f"- {os.path.join(eas_cache_dir, folder)}")
raise ValueError("Please download Fun model")
return model_name
def search_sub_dir_in_possible_folders(possible_folders, sub_dir_name="umt5-xxl"):
new_possible_folders = []
# Check if the model exists in any of the possible folders within folder_paths.models_dir
for folder in possible_folders:
candidate_path = os.path.join(folder_paths.models_dir, folder)
if os.path.exists(candidate_path) and os.path.isdir(candidate_path):
new_possible_folders.append(candidate_path)
for sub_dir in os.listdir(candidate_path):
new_possible_folders.append(os.path.join(candidate_path, sub_dir))
# If model_name is still None, check eas_cache_dir for each possible folder
if os.path.exists(eas_cache_dir):
for folder in possible_folders:
candidate_path = os.path.join(eas_cache_dir, folder)
if os.path.exists(candidate_path) and os.path.isdir(candidate_path):
new_possible_folders.append(candidate_path)
for sub_dir in os.listdir(candidate_path):
new_possible_folders.append(os.path.join(candidate_path, sub_dir))
for folder in new_possible_folders + possible_folders:
final_possible_folder = os.path.join(folder, sub_dir_name)
final_possible_folder_basename = os.path.join(folder, os.path.basename(sub_dir_name))
if os.path.exists(final_possible_folder) and os.path.isdir(final_possible_folder):
return final_possible_folder
if os.path.exists(final_possible_folder_basename) and os.path.isdir(final_possible_folder_basename):
return final_possible_folder_basename
print(f"Please download {sub_dir_name} tokenizer model to one of the following directories:")
for folder in possible_folders:
print(f"- {os.path.join(folder_paths.models_dir, folder)}")
if os.path.exists(eas_cache_dir):
print(f"- {os.path.join(eas_cache_dir, folder)}")
raise ValueError("Please download Fun model")
|