# utils.py
from enum import Enum, auto
import torch
from huggingface_hub import hf_hub_download
from PIL import Image, ImageEnhance, ImageFilter
import cv2
import numpy as np
from refiners.fluxion.utils import load_from_safetensors, tensor_to_image
from refiners.foundationals.clip import CLIPTextEncoderL
from refiners.foundationals.latent_diffusion import SD1UNet
from refiners.foundationals.latent_diffusion.stable_diffusion_1 import SD1Autoencoder
from refiners.foundationals.latent_diffusion.stable_diffusion_1.ic_light import ICLight
def load_ic_light(device: torch.device, dtype: torch.dtype) -> ICLight:
return ICLight(
patch_weights=load_from_safetensors(
path=hf_hub_download(
repo_id="refiners/sd15.ic_light.fc",
filename="model.safetensors",
revision="ea10b4403e97c786a98afdcbdf0e0fec794ea542",
),
),
unet=SD1UNet(in_channels=4, device=device, dtype=dtype).load_from_safetensors(
tensors_path=hf_hub_download(
repo_id="refiners/sd15.realistic_vision.v5_1.unet",
filename="model.safetensors",
revision="94f74be7adfd27bee330ea1071481c0254c29989",
)
),
clip_text_encoder=CLIPTextEncoderL(device=device, dtype=dtype).load_from_safetensors(
tensors_path=hf_hub_download(
repo_id="refiners/sd15.realistic_vision.v5_1.text_encoder",
filename="model.safetensors",
revision="7f6fa1e870c8f197d34488e14b89e63fb8d7fd6e",
)
),
lda=SD1Autoencoder(device=device, dtype=dtype).load_from_safetensors(
tensors_path=hf_hub_download(
repo_id="refiners/sd15.realistic_vision.v5_1.autoencoder",
filename="model.safetensors",
revision="99f089787a6e1a852a0992da1e286a19fcbbaa50",
)
),
device=device,
dtype=dtype,
)
def resize_modulo_8(
image: Image.Image,
size: int = 768,
resample: Image.Resampling | None = None,
on_short: bool = True,
) -> Image.Image:
"""이미지 크기를 8의 배수로 조정"""
assert size % 8 == 0, "Size must be a multiple of 8 because this is the latent compression size."
side_size = min(image.size) if on_short else max(image.size)
scale = size / (side_size * 8)
new_size = (int(image.width * scale) * 8, int(image.height * scale) * 8)
return image.resize(new_size, resample=resample or Image.Resampling.LANCZOS)
def adjust_image(
image: Image.Image,
brightness=0.0,
contrast=0.0,
temperature=0.0,
saturation=0.0,
tint=0.0,
blur_intensity=0,
exposure=0.0,
vibrance=0.0,
color_mixer_blues=0.0,
) -> Image.Image:
"""이미지 조정 함수"""
image = image.convert('RGB')
# 노출 조정 (Exposure)
if exposure != 0.0:
# Exposure ranges from -5 to 5, where 0 is neutral
exposure_factor = 1 + (exposure / 5.0)
exposure_factor = max(exposure_factor, 0.01) # Prevent zero or negative
enhancer = ImageEnhance.Brightness(image)
image = enhancer.enhance(exposure_factor)
# 밝기 조정
if brightness != 0.0:
# Brightness ranges from -5 to 5, mapped to brightness factor
brightness_factor = 1 + (brightness / 5.0)
brightness_factor = max(brightness_factor, 0.01) # Prevent zero or negative
enhancer = ImageEnhance.Brightness(image)
image = enhancer.enhance(brightness_factor)
# 대비 조정
if contrast != 0.0:
# Contrast ranges from -100 to 100, mapped to contrast factor
contrast_factor = 1 + (contrast / 100.0)
contrast_factor = max(contrast_factor, 0.01) # Prevent zero or negative
enhancer = ImageEnhance.Contrast(image)
image = enhancer.enhance(contrast_factor)
# 채도 조정 (Vibrance)
if vibrance != 0.0:
# Vibrance simulates adjusting the saturation; positive increases saturation, negative decreases
vibrance_factor = 1 + (vibrance / 100.0)
vibrance_factor = max(vibrance_factor, 0.0) # Prevent negative saturation
enhancer = ImageEnhance.Color(image)
image = enhancer.enhance(vibrance_factor)
# 채도 조정 (Saturation)
if saturation != 0.0:
# Saturation ranges from -100 to 100, mapped to saturation factor
saturation_factor = 1 + (saturation / 100.0)
saturation_factor = max(saturation_factor, 0.0) # Prevent negative saturation
enhancer = ImageEnhance.Color(image)
image = enhancer.enhance(saturation_factor)
# 색온도 조정
if temperature != 0.0:
# To prevent division by zero, adjust temperature calculation
temp_factor = 1 + (temperature / 100.0)
temp_factor = max(temp_factor, 0.01) # Prevent zero or negative
r, g, b = image.split()
r = r.point(lambda i: i * temp_factor)
b = b.point(lambda i: i / temp_factor)
image = Image.merge('RGB', (r, g, b))
# 색조 조정 (Tint)
if tint != 0.0:
image_np = np.array(image)
image_hsv = cv2.cvtColor(image_np, cv2.COLOR_RGB2HSV).astype(np.float32)
image_hsv[:, :, 0] = (image_hsv[:, :, 0] + tint) % 180
image_hsv[:, :, 0] = np.clip(image_hsv[:, :, 0], 0, 179)
image_rgb = cv2.cvtColor(image_hsv.astype(np.uint8), cv2.COLOR_HSV2RGB)
image = Image.fromarray(image_rgb)
# 블러 적용
if blur_intensity > 0:
image = image.filter(ImageFilter.GaussianBlur(radius=blur_intensity))
# Color Mixer (Blues)
if color_mixer_blues != 0.0:
image_np = np.array(image).astype(np.float32)
# Adjust the blue channel
image_np[:, :, 2] = np.clip(image_np[:, :, 2] + (color_mixer_blues / 100.0) * 255, 0, 255)
image = Image.fromarray(image_np.astype(np.uint8))
return image
class LightingPreference(str, Enum):
LEFT = auto()
RIGHT = auto()
TOP = auto()
BOTTOM = auto()
NONE = auto()
def get_init_image(self, width: int, height: int, interval: tuple[float, float] = (0.0, 1.0)) -> Image.Image | None:
"""조명 선호도에 따른 그라데이션 이미지 생성"""
start, end = interval
match self:
case LightingPreference.LEFT:
tensor = torch.linspace(end, start, width).repeat(1, 1, height, 1)
case LightingPreference.RIGHT:
tensor = torch.linspace(start, end, width).repeat(1, 1, height, 1)
case LightingPreference.TOP:
tensor = torch.linspace(end, start, height).repeat(1, 1, width, 1).transpose(2, 3)
case LightingPreference.BOTTOM:
tensor = torch.linspace(start, end, height).repeat(1, 1, width, 1).transpose(2, 3)
case LightingPreference.NONE:
return None
return tensor_to_image(tensor).convert("RGB")
@classmethod
def from_str(cls, value: str):
match value.lower():
case "left":
return LightingPreference.LEFT
case "right":
return LightingPreference.RIGHT
case "top":
return LightingPreference.TOP
case "bottom":
return LightingPreference.BOTTOM
case "none":
return LightingPreference.NONE
case _:
raise ValueError(f"Invalid lighting preference: {value}")