120 lines
3.8 KiB
Python
120 lines
3.8 KiB
Python
import os
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from typing import Union, List, Optional
|
|
|
|
def get_model_path(model_name: str) -> str:
|
|
"""
|
|
Get the path to a model file in the models directory.
|
|
|
|
Args:
|
|
model_name: Name of the model file or directory
|
|
|
|
Returns:
|
|
Full path to the model
|
|
"""
|
|
# Determine the path to the models directory
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
models_dir = os.path.join(current_dir, "models")
|
|
|
|
# Check if the model exists
|
|
model_path = os.path.join(models_dir, model_name)
|
|
if not os.path.exists(model_path):
|
|
raise FileNotFoundError(f"Model {model_name} not found in {models_dir}")
|
|
|
|
return model_path
|
|
|
|
def get_lora_path(lora_name: str) -> str:
|
|
"""
|
|
Get the path to a LoRA file in the loras directory.
|
|
|
|
Args:
|
|
lora_name: Name of the LoRA file
|
|
|
|
Returns:
|
|
Full path to the LoRA file
|
|
"""
|
|
# Determine the path to the loras directory
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
loras_dir = os.path.join(current_dir, "loras")
|
|
|
|
# Check if the LoRA file exists
|
|
lora_path = os.path.join(loras_dir, lora_name)
|
|
if not os.path.exists(lora_path):
|
|
raise FileNotFoundError(f"LoRA file {lora_name} not found in {loras_dir}")
|
|
|
|
return lora_path
|
|
|
|
def comfy_to_pil(image: torch.Tensor) -> Image.Image:
|
|
"""
|
|
Convert a ComfyUI image tensor to a PIL Image.
|
|
|
|
Args:
|
|
image: ComfyUI image tensor (1, H, W, 3) in range [0, 1]
|
|
|
|
Returns:
|
|
PIL Image
|
|
"""
|
|
# Convert to numpy array and scale to [0, 255]
|
|
image_np = np.clip(255. * image[0].cpu().numpy(), 0, 255).astype(np.uint8)
|
|
# Convert to PIL Image
|
|
return Image.fromarray(image_np)
|
|
|
|
def pil_to_comfy(image: Union[Image.Image, List[Image.Image], None]) -> Optional[torch.Tensor]:
|
|
"""
|
|
Convert a PIL Image or list of PIL Images to a ComfyUI image tensor.
|
|
|
|
Args:
|
|
image: PIL Image, list of PIL Images, or None
|
|
|
|
Returns:
|
|
ComfyUI image tensor (1, H, W, 3) in range [0, 1] or None if input is None
|
|
"""
|
|
if image is None:
|
|
return None
|
|
|
|
# Handle list of PIL images - take the first one
|
|
if isinstance(image, list):
|
|
if len(image) == 0:
|
|
return None
|
|
image = image[0] # Take the first image from the list
|
|
|
|
# Convert to numpy array and scale to [0, 1]
|
|
image_np = np.array(image).astype(np.float32) / 255.0
|
|
|
|
# Ensure the image has the right shape (H, W, 3)
|
|
if len(image_np.shape) == 2: # Grayscale image
|
|
image_np = np.stack([image_np, image_np, image_np], axis=-1)
|
|
elif image_np.shape[-1] == 4: # RGBA image
|
|
image_np = image_np[..., :3] # Remove alpha channel
|
|
|
|
# Convert to torch tensor and add batch dimension
|
|
return torch.from_numpy(image_np)[None,]
|
|
|
|
def center_crop_and_resize(image: Union[torch.Tensor, Image.Image], target_size: int = 512) -> Union[torch.Tensor, Image.Image]:
|
|
"""
|
|
Center crop and resize an image.
|
|
|
|
Args:
|
|
image: Image to process (ComfyUI tensor or PIL Image)
|
|
target_size: Target size for width and height
|
|
|
|
Returns:
|
|
Processed image in the same format as input
|
|
"""
|
|
# Handle ComfyUI tensor
|
|
if isinstance(image, torch.Tensor):
|
|
pil_image = comfy_to_pil(image)
|
|
result = center_crop_and_resize(pil_image, target_size)
|
|
return pil_to_comfy(result)
|
|
|
|
# Handle PIL Image
|
|
w, h = image.size
|
|
min_size = min(w, h)
|
|
cropped = image.crop(((w - min_size) // 2,
|
|
(h - min_size) // 2,
|
|
(w + min_size) // 2,
|
|
(h + min_size) // 2))
|
|
resized = cropped.resize((target_size, target_size), Image.LANCZOS)
|
|
return resized |