122 lines
4.5 KiB
Python
122 lines
4.5 KiB
Python
import os
|
|
import json
|
|
import torch
|
|
from safetensors import safe_open
|
|
|
|
def load_diffusers_model(model_path, output_vae=True, output_clip=True):
|
|
try:
|
|
model = load_model_from_path(model_path)
|
|
vae = load_vae_from_path(model_path) if output_vae else None
|
|
clip = load_text_encoder_from_path(model_path, encoder_id=None) if output_clip else None
|
|
return model, vae, clip
|
|
except Exception as e:
|
|
print(f"Error loading diffusers model from {model_path}: {e}")
|
|
raise
|
|
|
|
def load_diffusers_sdxl_model(model_path, output_vae=True, output_clip=True, embedding_directory=None):
|
|
try:
|
|
model = load_model_from_path(model_path)
|
|
vae = load_vae_from_path(model_path) if output_vae else None
|
|
clip = load_text_encoder_from_path(model_path, encoder_id=1) if output_clip else None
|
|
text_encoder_2 = load_text_encoder_from_path(model_path, encoder_id=2)
|
|
|
|
if embedding_directory:
|
|
# Load or handle embeddings if applicable
|
|
pass
|
|
|
|
return model, vae, clip, text_encoder_2
|
|
except Exception as e:
|
|
print(f"Error loading SDXL diffusers model from {model_path}: {e}")
|
|
raise
|
|
|
|
def load_model_from_path(model_path):
|
|
try:
|
|
print(f"Loading model from {model_path}")
|
|
check_file_permissions(model_path)
|
|
if model_path.endswith('.safetensors'):
|
|
model = load_safetensors(model_path)
|
|
else:
|
|
model = load_json(model_path)
|
|
return model
|
|
except PermissionError as e:
|
|
print(f"Permission denied when accessing {model_path}. Please check the file permissions.")
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error loading model from {model_path}: {e}")
|
|
raise
|
|
|
|
def load_vae_from_path(model_path):
|
|
try:
|
|
print(f"Loading VAE from {model_path}")
|
|
vae_path = os.path.join(model_path, "vae", "diffusion_pytorch_model.safetensors")
|
|
check_file_permissions(vae_path)
|
|
vae = load_safetensors(vae_path)
|
|
return vae
|
|
except PermissionError as e:
|
|
print(f"Permission denied when accessing {vae_path}. Please check the file permissions.")
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error loading VAE from {vae_path}: {e}")
|
|
raise
|
|
|
|
def load_clip_from_path(model_path):
|
|
try:
|
|
print(f"Loading CLIP model from {model_path}")
|
|
clip_path = os.path.join(model_path, "clip", "diffusion_pytorch_model.safetensors")
|
|
check_file_permissions(clip_path)
|
|
clip = load_safetensors(clip_path)
|
|
return clip
|
|
except PermissionError as e:
|
|
print(f"Permission denied when accessing {clip_path}. Please check the file permissions.")
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error loading CLIP model from {clip_path}: {e}")
|
|
raise
|
|
|
|
def load_text_encoder_from_path(model_path, encoder_id=None):
|
|
try:
|
|
if encoder_id is not None:
|
|
encoder_path = os.path.join(model_path, f"text_encoder_{encoder_id}", "model.safetensors")
|
|
print(f"Loading text encoder {encoder_id} from {encoder_path}")
|
|
else:
|
|
encoder_path = os.path.join(model_path, "text_encoder", "model.safetensors")
|
|
print(f"Loading text encoder from {encoder_path}")
|
|
|
|
check_file_permissions(encoder_path)
|
|
text_encoder = load_safetensors(encoder_path)
|
|
return text_encoder
|
|
except PermissionError as e:
|
|
print(f"Permission denied when accessing {encoder_path}. Please check the file permissions.")
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error loading text encoder from {encoder_path}: {e}")
|
|
raise
|
|
|
|
def load_json(json_path):
|
|
try:
|
|
print(f"Loading JSON from {json_path}")
|
|
with open(json_path, 'r') as f:
|
|
return json.load(f)
|
|
except Exception as e:
|
|
print(f"Error loading JSON from {json_path}: {e}")
|
|
raise
|
|
|
|
def load_safetensors(file_path):
|
|
try:
|
|
print(f"Loading safetensors from {file_path}")
|
|
with safe_open(file_path, framework="pt", device="cpu") as f:
|
|
return {k: f.get_tensor(k) for k in f.keys()}
|
|
except Exception as e:
|
|
print(f"Error loading safetensors from {file_path}: {e}")
|
|
raise
|
|
|
|
def check_file_permissions(file_path):
|
|
try:
|
|
if os.access(file_path, os.R_OK | os.W_OK):
|
|
print(f"Read and write permissions are available for {file_path}")
|
|
else:
|
|
raise PermissionError(f"Permission issue with {file_path}")
|
|
except Exception as e:
|
|
print(f"Error checking permissions for {file_path}: {e}")
|
|
raise
|