705 lines
28 KiB
Python
705 lines
28 KiB
Python
"""
|
||
Copyright (c) 2026 Quentin Lengele
|
||
"""
|
||
import os
|
||
import sys
|
||
import torch
|
||
import cv2
|
||
import random
|
||
import numpy as np
|
||
from pathlib import Path
|
||
from PIL import Image, ImageFilter
|
||
from scipy.ndimage import gaussian_filter
|
||
|
||
extension_folder = os.path.dirname(os.path.realpath(__file__))
|
||
|
||
parent_folder = os.path.dirname(extension_folder)
|
||
if parent_folder not in sys.path:
|
||
sys.path.append(parent_folder)
|
||
|
||
from .matgen import imgops as ops
|
||
from .matgen import architecture as arch
|
||
|
||
import folder_paths
|
||
|
||
# =====================================================================================================================
|
||
|
||
def load_model(model_path, device):
|
||
state_dict = torch.load(model_path, map_location=device)
|
||
model = arch.RRDB_Net(3, 3, 32, 12, gc=32, upscale=1, norm_type=None, act_type='leakyrelu',
|
||
mode='CNA', res_scale=1, upsample_mode='upconv')
|
||
model.load_state_dict(state_dict, strict=True)
|
||
model.eval()
|
||
for v in model.parameters():
|
||
v.requires_grad = False
|
||
return model.to(device)
|
||
|
||
|
||
def process(img, model, device):
|
||
"""
|
||
Bridge ComfyUI (RGB Float) to Model (BGR Logic)
|
||
"""
|
||
|
||
# --- Start of Original Library Logic ---
|
||
# Normalize back to float based on the uint8 max (255)
|
||
img = img * 1. / np.iinfo(img.dtype).max
|
||
|
||
# Model expects BGR input, but original code flips it again
|
||
# to [2,1,0] before feeding to model. We follow that logic:
|
||
img = img[:, :, [2, 1, 0]]
|
||
img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()
|
||
img_LR = img.unsqueeze(0).to(device)
|
||
|
||
with torch.no_grad():
|
||
output = model(img_LR).data.squeeze(0).float().cpu().clamp_(0, 1).numpy()
|
||
|
||
# Flip channels back
|
||
output = output[[2, 1, 0], :, :]
|
||
output = np.transpose(output, (1, 2, 0))
|
||
output = (output * 255.0).round()
|
||
# --- End of Original Library Logic ---
|
||
|
||
return output # Returns 0-255 BGR NumPy array
|
||
|
||
# =====================================================================================================================
|
||
|
||
def _validate_bhwc_image(image: torch.Tensor) -> None:
|
||
if image.dim() != 4:
|
||
raise ValueError("Expected IMAGE tensor in BHWC format")
|
||
if image.shape[3] not in (3, 4):
|
||
raise ValueError(f"Expected 3 or 4 channels, got {image.shape[3]}")
|
||
|
||
|
||
def _clamp01(x: torch.Tensor) -> torch.Tensor:
|
||
return torch.clamp(x, 0.0, 1.0)
|
||
|
||
|
||
def _to_chw(img_hwc: torch.Tensor) -> torch.Tensor:
|
||
"""Convert HWC to CHW format."""
|
||
return img_hwc.permute(2, 0, 1)
|
||
|
||
|
||
def _to_hwc(img_chw: torch.Tensor) -> torch.Tensor:
|
||
"""Convert CHW to HWC format."""
|
||
return img_chw.permute(1, 2, 0)
|
||
|
||
|
||
def _make_tile_preview(img: torch.Tensor, tiles_x: int, tiles_y: int) -> torch.Tensor:
|
||
"""Create a tiled preview of the texture. img: CHW format."""
|
||
tiles_x = max(1, int(tiles_x))
|
||
tiles_y = max(1, int(tiles_y))
|
||
return torch.tile(img, (1, tiles_y, tiles_x))
|
||
|
||
# =====================================================================================================================
|
||
|
||
def adjust_normal_strength(normal_bgr, strength=1.0, invert_y=False):
|
||
"""
|
||
Scales the intensity of the normal map vectors and optionally inverts the Y channel.
|
||
"""
|
||
# Convert 0-255 to -1.0 to 1.0
|
||
norm = normal_bgr.astype(np.float32) / 127.5 - 1.0
|
||
|
||
# Scale X and Y
|
||
norm[:, :, 2] *= strength # R (X)
|
||
norm[:, :, 1] *= strength # G (Y)
|
||
|
||
# UNREAL ENGINE FIX: Invert the Green channel
|
||
if invert_y:
|
||
norm[:, :, 1] *= -1.0
|
||
|
||
# Re-normalize vectors
|
||
mag = np.sqrt(np.sum(np.square(norm), axis=2, keepdims=True))
|
||
norm = norm / (mag + 1e-6)
|
||
|
||
return ((norm + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
|
||
|
||
|
||
def generate_ao_map(displacement_arr, blur_radius=2):
|
||
# Invert height: high areas = less occluded
|
||
inv_height = 1.0 - (displacement_arr.astype(np.float32) / 255.0)
|
||
inv_height_img = Image.fromarray((inv_height * 255).astype(np.uint8))
|
||
ao_img = inv_height_img.filter(ImageFilter.GaussianBlur(radius=blur_radius))
|
||
return np.array(ao_img)
|
||
|
||
|
||
def generate_metallic_map(albedo_bgr, sat_thresh, min_val):
|
||
"""Internal helper for HSV metallic detection"""
|
||
hsv = cv2.cvtColor(albedo_bgr, cv2.COLOR_BGR2HSV).astype(np.float32)
|
||
sat = hsv[:, :, 1] / 255.0
|
||
val = hsv[:, :, 2] / 255.0
|
||
|
||
# Neutral metals (Low saturation, high enough brightness)
|
||
metal_mask = np.where((sat < sat_thresh) & (val > min_val), 1.0, 0.0)
|
||
|
||
# Add slight detail based on luminance contrast
|
||
gray = cv2.cvtColor(albedo_bgr, cv2.COLOR_BGR2GRAY)
|
||
laplacian = cv2.convertScaleAbs(cv2.Laplacian(gray, cv2.CV_32F)) / 255.0
|
||
metal_mask = np.clip(metal_mask + (laplacian * 0.1), 0, 1)
|
||
|
||
return (gaussian_filter(metal_mask, sigma=0.5) * 255).astype(np.uint8)
|
||
|
||
|
||
def prepare_for_merge(img_arr, target_w, target_h):
|
||
# Ensure 8-bit uint8
|
||
img_arr = img_arr.astype(np.uint8)
|
||
# Ensure dimensions match model output
|
||
if (img_arr.shape[0] != target_h) or (img_arr.shape[1] != target_w):
|
||
img_arr = cv2.resize(img_arr, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
|
||
return img_arr
|
||
|
||
|
||
def reprocess_map(img_array, min_val=0.0, max_val=1.0):
|
||
"""
|
||
Remaps a 0-255 grayscale map to a specific min/max range.
|
||
min_val/max_val are expected as 0.0 to 1.0 floats.
|
||
"""
|
||
# Convert to 0.0 - 1.0 float
|
||
img = img_array.astype(np.float32) / 255.0
|
||
|
||
# Linear remap: NewValue = min + (Value * (max - min))
|
||
# This keeps the internal details/texture but compresses them into the range
|
||
img = min_val + (img * (max_val - min_val))
|
||
|
||
return (np.clip(img, 0, 1) * 255).astype(np.uint8)
|
||
|
||
|
||
def tensor_to_cv2(tensor):
|
||
"""
|
||
Standardizes ComfyUI BHWC tensors to OpenCV HWC BGR/Gray arrays.
|
||
"""
|
||
# 1. Detach and move to CPU
|
||
img = tensor.detach().cpu()
|
||
|
||
# 2. Handle Batch [B, H, W, C] -> [H, W, C]
|
||
if img.ndim == 4:
|
||
img = img[0]
|
||
|
||
# 3. Handle Channels
|
||
if img.shape[-1] == 3:
|
||
# If it's [H, W, 3] (RGB), convert to BGR for OpenCV
|
||
img_np = img.numpy()[:, :, ::-1]
|
||
elif img.shape[-1] == 1:
|
||
# If it's [H, W, 1], squeeze to [H, W] (Grayscale)
|
||
img_np = img.squeeze(-1).numpy()
|
||
else:
|
||
img_np = img.numpy()
|
||
|
||
# 4. Final scaling and casting
|
||
return (img_np * 255.0).clip(0, 255).astype(np.uint8)
|
||
|
||
|
||
def cv2_to_tensor(img):
|
||
"""
|
||
Converts a NumPy/OpenCV image (HWC, BGR or Gray)
|
||
to a ComfyUI-standard Tensor (BHWC, RGB).
|
||
"""
|
||
# 1. Handle Grayscale vs BGR
|
||
if len(img.shape) == 2: # Grayscale (H, W)
|
||
# Convert to (H, W, 1) then repeat to (H, W, 3)
|
||
# because ComfyUI save nodes prefer 3-channel RGB
|
||
img = img[:, :, np.newaxis]
|
||
img = np.repeat(img, 3, axis=2)
|
||
elif img.shape[2] == 3:
|
||
# BGR to RGB (HWC)
|
||
img = cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_BGR2RGB)
|
||
|
||
# 2. Convert to Tensor and normalize
|
||
# Standard ComfyUI format is [Batch, Height, Width, Channels]
|
||
# We do NOT use .permute(2, 0, 1) here because Comfy wants HWC
|
||
tensor = torch.from_numpy(img).float().div(255.0)
|
||
|
||
# 3. Add Batch Dimension -> [1, H, W, 3]
|
||
return tensor.unsqueeze(0)
|
||
|
||
|
||
# =====================================================================================================================
|
||
# =====================================================================================================================
|
||
# =====================================================================================================================
|
||
|
||
class QFXPBRGenerator:
|
||
def __init__(self):
|
||
self.output_dir = folder_paths.get_output_directory()
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE", {"forceInput": True}),
|
||
|
||
"texture_name": ("STRING", {"default": "MyTexture"}),
|
||
|
||
"seamless_mode": (["off", "seamless", "mirror", "replicate"], {
|
||
"default": "seamless",
|
||
"tooltip": ""
|
||
}),
|
||
|
||
# --- Normal ---
|
||
"normal_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.05}),
|
||
"normal_invert_y": ("BOOLEAN", {"default": True}),
|
||
# --- Roughness Remapping ---
|
||
"roughness_min": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"roughness_max": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
# --- Metallic Remapping ---
|
||
"metallic_min": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"metallic_max": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"metallic_threshold": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
# --- AO & Displacement ---
|
||
"ao_blur": ("INT", {"default": 2, "min": 0, "max": 10, "step": 1}),
|
||
"ao_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05}),
|
||
# ---
|
||
"save_textures": ("BOOLEAN", {"default": True}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE")
|
||
RETURN_NAMES = ("BaseColor", "Normal", "Roughness", "Displacement", "AO", "Metallic", "ORM")
|
||
FUNCTION = "generate"
|
||
OUTPUT_NODE = True
|
||
CATEGORY = "QFX"
|
||
|
||
def generate(self, image, texture_name, seamless_mode, normal_strength, normal_invert_y, roughness_min,
|
||
roughness_max, metallic_min, metallic_max, ao_blur,
|
||
ao_strength, metallic_threshold, save_textures):
|
||
|
||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
base_path = Path(__file__).parent.resolve()
|
||
|
||
# 1. Setup Models
|
||
normal_model_path = str(base_path / 'models' / '1x_NormalMapGenerator-CX-Lite_200000_G.pth')
|
||
other_model_path = str(base_path / 'models' / '1x_FrankenMapGenerator-CX-Lite_215000_G.pth')
|
||
model_normal = load_model(normal_model_path, device)
|
||
model_other = load_model(other_model_path, device)
|
||
|
||
# 2. Get Input Image
|
||
img_input = image[0]
|
||
albedo_bgr = (img_input.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)[:, :, [2, 1, 0]]
|
||
|
||
albedo_input = albedo_bgr
|
||
if seamless_mode != "off":
|
||
albedo_input = ops.make_borders(albedo_bgr, seamless_mode, 16)
|
||
|
||
# 3. Generate AI Maps
|
||
raw_normal_bgr = process(albedo_input, model_normal, device)
|
||
other_maps_bgr = process(albedo_input, model_other, device)
|
||
|
||
if seamless_mode != "off":
|
||
raw_normal_bgr = ops.crop_seamless(raw_normal_bgr, 16)
|
||
other_maps_bgr = ops.crop_seamless(other_maps_bgr, 16)
|
||
|
||
# 4. Refine with Min/Max Remapping
|
||
|
||
# Normal
|
||
normal_map_bgr = adjust_normal_strength(raw_normal_bgr, normal_strength, normal_invert_y)
|
||
|
||
# Roughness (Remap AI output to user range)
|
||
roughness_gray = reprocess_map(other_maps_bgr[:, :, 1], roughness_min, roughness_max)
|
||
|
||
# Displacement
|
||
displacement_gray = other_maps_bgr[:, :, 0]
|
||
|
||
# AO (Strength here acts as a darkness multiplier for the occlusion)
|
||
ao_raw = generate_ao_map(displacement_gray, blur_radius=ao_blur)
|
||
ao_gray = reprocess_map(ao_raw, min_val=(1.0 - ao_strength), max_val=1.0)
|
||
|
||
# Metallic (Heuristic detection -> then remap intensity)
|
||
raw_metallic = generate_metallic_map(albedo_bgr, sat_thresh=metallic_threshold, min_val=0.2)
|
||
metallic_gray = reprocess_map(raw_metallic, metallic_min, metallic_max)
|
||
|
||
# 5. ORM Packing & Return
|
||
target_h, target_w = roughness_gray.shape[:2]
|
||
r_ao = prepare_for_merge(ao_gray, target_w, target_h)
|
||
g_rough = prepare_for_merge(roughness_gray, target_w, target_h)
|
||
b_metal = prepare_for_merge(metallic_gray, target_w, target_h)
|
||
orm_bgr_packed = cv2.merge([b_metal, g_rough, r_ao])
|
||
|
||
if save_textures:
|
||
# Saving logic (uses the refined maps)
|
||
self._save_all(texture_name, albedo_bgr, normal_map_bgr, roughness_gray,
|
||
displacement_gray, ao_gray, metallic_gray, orm_bgr_packed)
|
||
|
||
return (
|
||
cv2_to_tensor(albedo_bgr),
|
||
cv2_to_tensor(normal_map_bgr),
|
||
cv2_to_tensor(roughness_gray),
|
||
cv2_to_tensor(displacement_gray),
|
||
cv2_to_tensor(ao_gray),
|
||
cv2_to_tensor(metallic_gray),
|
||
cv2_to_tensor(orm_bgr_packed)
|
||
)
|
||
|
||
def _save_all(self, name, albedo, normal, rough, disp, ao, metal, orm):
|
||
maps = {
|
||
"BaseColor": albedo, "Normal": normal, "Roughness": rough,
|
||
"Displacement": disp, "AO": ao, "Metallic": metal, "ORM": orm
|
||
}
|
||
for suffix, data in maps.items():
|
||
cv2.imwrite(os.path.join(self.output_dir, f"{name}_{suffix}.png"), data)
|
||
|
||
# =====================================================================================================================
|
||
# =====================================================================================================================
|
||
# =====================================================================================================================
|
||
|
||
class QFXPBRGeneratorAdvanced:
|
||
def __init__(self):
|
||
self.output_dir = folder_paths.get_output_directory()
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
|
||
"image": ("IMAGE",),
|
||
|
||
"texture_name": ("STRING", {"default": "MyTexture"}),
|
||
|
||
"seamless_mode": (["off", "seamless", "mirror", "replicate"], {
|
||
"default": "seamless",
|
||
"tooltip": ""
|
||
}),
|
||
|
||
# Marigold outputs
|
||
"marigold_appearance": ("IMAGE",),
|
||
"marigold_lighting": ("IMAGE",),
|
||
|
||
# Albedo choice
|
||
"albedo_source": (["appearance", "lighting"], {
|
||
"default": "lighting",
|
||
"tooltip": "Which Marigold output to use for albedo"
|
||
}),
|
||
|
||
# Gamma correction - separate controls
|
||
"gamma_albedo": ("FLOAT", {
|
||
"default": 0.45,
|
||
"min": 0.1,
|
||
"max": 3.0,
|
||
"step": 0.01,
|
||
"display": "number",
|
||
"tooltip": "Gamma correction for albedo (applied to both appearance and lighting sources)"
|
||
}),
|
||
"gamma_metal_rough": ("FLOAT", {
|
||
"default": 2.2,
|
||
"min": 0.1,
|
||
"max": 3.0,
|
||
"step": 0.01,
|
||
"display": "number",
|
||
"tooltip": "Gamma correction for metallic and roughness maps"
|
||
}),
|
||
"gamma_lighting_ao": ("FLOAT", {
|
||
"default": 0.45,
|
||
"min": 0.1,
|
||
"max": 3.0,
|
||
"step": 0.01,
|
||
"display": "number",
|
||
"tooltip": "Gamma correction for lighting and AO maps"
|
||
}),
|
||
|
||
# --- Roughness Remapping ---
|
||
"roughness_min": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"roughness_max": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
|
||
# --- Metallic Remapping ---
|
||
"metallic_min": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"metallic_max": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
|
||
# --- Normal ---
|
||
"normal_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.05}),
|
||
"normal_invert_y": ("BOOLEAN", {"default": True}),
|
||
|
||
# --- AO & Displacement ---
|
||
"ao_blur": ("INT", {"default": 2, "min": 0, "max": 10, "step": 1}),
|
||
"ao_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05}),
|
||
|
||
# ---
|
||
"save_textures": ("BOOLEAN", {"default": True}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE")
|
||
RETURN_NAMES = ("Original", "BaseColor", "Normal", "Roughness", "Displacement", "AO", "Metallic", "ORM")
|
||
FUNCTION = "generate"
|
||
OUTPUT_NODE = True
|
||
CATEGORY = "QFX"
|
||
|
||
def generate(self, image, texture_name, seamless_mode,
|
||
marigold_appearance, marigold_lighting, albedo_source,
|
||
gamma_albedo, gamma_metal_rough, gamma_lighting_ao,
|
||
roughness_min, roughness_max,
|
||
metallic_min, metallic_max,
|
||
normal_strength, normal_invert_y,
|
||
ao_blur, ao_strength,
|
||
save_textures):
|
||
|
||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
base_path = Path(__file__).parent.resolve()
|
||
|
||
# 1. Extract Maps from Marigold Tensors
|
||
# These return BHWC tensors
|
||
delighted_albedo, roughness, metallic, ao = self.extract(marigold_appearance,
|
||
marigold_lighting,
|
||
albedo_source,
|
||
gamma_albedo,
|
||
gamma_metal_rough,
|
||
gamma_lighting_ao)
|
||
|
||
if delighted_albedo.shape[1] == 3: # If it's [B, C, H, W]
|
||
delighted_albedo = delighted_albedo.permute(0, 2, 3, 1)
|
||
|
||
# Convert to CV2 for internal processing (Grayscale/BGR)
|
||
delighted_albedo_bgr = tensor_to_cv2(delighted_albedo)
|
||
# roughness_gray = tensor_to_cv2(roughness)
|
||
metallic_gray = tensor_to_cv2(metallic)
|
||
|
||
# roughness_gray = reprocess_map(roughness_gray, roughness_min, roughness_max)
|
||
metallic_gray = reprocess_map(metallic_gray, metallic_min, metallic_max)
|
||
|
||
# Setup AI Models
|
||
normal_model_path = str(base_path / 'models' / '1x_NormalMapGenerator-CX-Lite_200000_G.pth')
|
||
other_model_path = str(base_path / 'models' / '1x_FrankenMapGenerator-CX-Lite_215000_G.pth')
|
||
model_normal = load_model(normal_model_path, device)
|
||
model_other = load_model(other_model_path, device)
|
||
|
||
# Generate AI Maps
|
||
# Note: 'process' expects a single image tensor.
|
||
img_input = image[0]
|
||
albedo_bgr = (img_input.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)[:, :, [2, 1, 0]]
|
||
|
||
albedo_input = albedo_bgr
|
||
if seamless_mode != "off":
|
||
albedo_input = ops.make_borders(albedo_bgr, seamless_mode, 16)
|
||
|
||
raw_normal_bgr = process(albedo_input, model_normal, device)
|
||
other_maps_bgr = process(albedo_input, model_other, device)
|
||
|
||
if seamless_mode != "off":
|
||
raw_normal_bgr = ops.crop_seamless(raw_normal_bgr, 16)
|
||
other_maps_bgr = ops.crop_seamless(other_maps_bgr, 16)
|
||
|
||
# Refine Normal
|
||
normal_map_bgr = adjust_normal_strength(raw_normal_bgr, normal_strength, normal_invert_y)
|
||
|
||
# Roughness
|
||
roughness_gray = reprocess_map(other_maps_bgr[:, :, 1], roughness_min, roughness_max)
|
||
|
||
# Displacement
|
||
displacement_gray = other_maps_bgr[:, :, 0]
|
||
|
||
# AO Generation
|
||
ao_raw = generate_ao_map(displacement_gray, blur_radius=ao_blur)
|
||
ao_gray = reprocess_map(ao_raw, min_val=(1.0 - ao_strength), max_val=1.0)
|
||
|
||
# ORM Packing
|
||
target_h, target_w = roughness_gray.shape[:2]
|
||
r_ao = prepare_for_merge(ao_gray, target_w, target_h)
|
||
g_rough = prepare_for_merge(roughness_gray, target_w, target_h)
|
||
b_metal = prepare_for_merge(metallic_gray, target_w, target_h)
|
||
|
||
# Ensure we are taking only ONE channel from each map for the ORM
|
||
r_ao_single = r_ao[:, :, 0] if r_ao.ndim == 3 else r_ao
|
||
g_rough_single = g_rough[:, :, 0] if g_rough.ndim == 3 else g_rough
|
||
b_metal_single = b_metal[:, :, 0] if b_metal.ndim == 3 else b_metal
|
||
|
||
# This creates a 3-channel BGR image
|
||
orm_bgr_packed = cv2.merge([b_metal_single, g_rough_single, r_ao_single])
|
||
|
||
# Save files
|
||
|
||
if save_textures:
|
||
self._save_all(texture_name, albedo_bgr, delighted_albedo_bgr, normal_map_bgr, roughness_gray,
|
||
displacement_gray, ao_gray, metallic_gray, orm_bgr_packed)
|
||
|
||
# -------
|
||
|
||
return (
|
||
cv2_to_tensor(albedo_bgr),
|
||
delighted_albedo,
|
||
cv2_to_tensor(normal_map_bgr),
|
||
cv2_to_tensor(roughness_gray),
|
||
cv2_to_tensor(displacement_gray),
|
||
cv2_to_tensor(ao_gray),
|
||
cv2_to_tensor(metallic_gray),
|
||
cv2_to_tensor(orm_bgr_packed)
|
||
)
|
||
|
||
def extract(self, marigold_appearance,
|
||
marigold_lighting,
|
||
albedo_source,
|
||
gamma_albedo,
|
||
gamma_metal_rough,
|
||
gamma_lighting_ao):
|
||
|
||
"""Extract PBR maps"""
|
||
|
||
print("\n" + "=" * 60)
|
||
print("PBR Extractor")
|
||
print("=" * 60)
|
||
print(f"Lighting shape: {marigold_lighting.shape}")
|
||
print(f"Albedo source: {albedo_source}")
|
||
print(f"Gamma - Albedo: {gamma_albedo}, Metal/Rough: {gamma_metal_rough}, Lighting/AO: {gamma_lighting_ao}")
|
||
|
||
# ===== ALBEDO =====
|
||
if albedo_source == "appearance":
|
||
# Get index 0 from appearance
|
||
albedo = marigold_appearance[0:1] if marigold_appearance.shape[0] > 0 else marigold_appearance
|
||
print("\n✓ Albedo from Appearance (index 0)")
|
||
else:
|
||
# Get index 0 from lighting
|
||
albedo = marigold_lighting[0:1] if marigold_lighting.shape[0] > 0 else marigold_lighting
|
||
print("\n✓ Albedo from Lighting (index 0)")
|
||
|
||
albedo = self.apply_gamma(albedo, gamma_albedo)
|
||
print(f" Gamma: {gamma_albedo}")
|
||
print(f" Range: [{albedo.min():.3f}, {albedo.max():.3f}]")
|
||
|
||
# ===== ROUGHNESS (Red channel from Appearance index 1) =====
|
||
if marigold_appearance.shape[0] >= 2:
|
||
material_batch = marigold_appearance[1:2]
|
||
roughness = material_batch[:, :, :, 0:1] # Red channel
|
||
roughness = self.channel_to_rgb(roughness)
|
||
roughness = self.apply_gamma(roughness, gamma_metal_rough)
|
||
print("\n✓ Roughness from Appearance index 1, RED channel")
|
||
print(f" Gamma: {gamma_metal_rough}")
|
||
print(f" Range: [{roughness.min():.3f}, {roughness.max():.3f}]")
|
||
else:
|
||
print("\n✗ WARNING: Appearance doesn't have index 1, using zeros")
|
||
roughness = torch.zeros_like(albedo)
|
||
|
||
# ===== METALLIC (Green channel from Appearance index 1) =====
|
||
if marigold_appearance.shape[0] >= 2:
|
||
material_batch = marigold_appearance[1:2]
|
||
metallic = material_batch[:, :, :, 1:2] # Green channel
|
||
metallic = self.channel_to_rgb(metallic)
|
||
metallic = self.apply_gamma(metallic, gamma_metal_rough)
|
||
print("\n✓ Metallic from Appearance index 1, GREEN channel")
|
||
print(f" Gamma: {gamma_metal_rough}")
|
||
print(f" Range: [{metallic.min():.3f}, {metallic.max():.3f}]")
|
||
else:
|
||
print("\n✗ WARNING: Appearance doesn't have index 1, using zeros")
|
||
metallic = torch.zeros_like(albedo)
|
||
|
||
# ===== LIGHTING / AO (Lighting index 1) =====
|
||
if marigold_lighting.shape[0] >= 2:
|
||
ao = marigold_lighting[1:2]
|
||
ao = self.grayscale(ao)
|
||
ao = self.apply_gamma(ao, gamma_lighting_ao)
|
||
print("\n✓ Lighting/AO from Lighting index 1")
|
||
print(f" Gamma: {gamma_lighting_ao}")
|
||
print(f" Range: [{ao.min():.3f}, {ao.max():.3f}]")
|
||
else:
|
||
print("\n✗ WARNING: Lighting doesn't have index 1, using albedo luminance")
|
||
ao = self.grayscale(albedo)
|
||
|
||
print("\n" + "=" * 60)
|
||
print("✓ Extraction Complete")
|
||
print("✓ PBR Pipeline created")
|
||
print("=" * 60 + "\n")
|
||
|
||
return albedo, roughness, metallic, ao
|
||
|
||
def apply_gamma(self, image, gamma):
|
||
"""Apply gamma correction"""
|
||
return torch.pow(torch.clamp(image, 0.0, 1.0), gamma)
|
||
|
||
def channel_to_rgb(self, channel):
|
||
"""Convert single channel to RGB"""
|
||
if channel.shape[-1] == 1:
|
||
return channel.repeat(1, 1, 1, 3)
|
||
return channel
|
||
|
||
def grayscale(self, image):
|
||
"""Convert to grayscale"""
|
||
if image.shape[-1] == 1:
|
||
return image
|
||
weights = torch.tensor([0.2989, 0.5870, 0.1140],
|
||
device=image.device, dtype=image.dtype)
|
||
gray = torch.sum(image * weights, dim=-1, keepdim=True)
|
||
return self.channel_to_rgb(gray)
|
||
|
||
def _save_all(self, name, original, albedo, normal, rough, disp, ao, metal, orm):
|
||
maps = {
|
||
"Original": original,
|
||
"BaseColor": albedo, "Normal": normal, "Roughness": rough,
|
||
"Displacement": disp, "AO": ao, "Metallic": metal, "ORM": orm
|
||
}
|
||
|
||
# Ensure the directory exists
|
||
if not os.path.exists(self.output_dir):
|
||
os.makedirs(self.output_dir, exist_ok=True)
|
||
|
||
print(f"--- PBR Saver: Exporting to {self.output_dir} ---")
|
||
|
||
for suffix, data in maps.items():
|
||
# 1. Force convert to NumPy if it's a Tensor
|
||
if torch.is_tensor(data):
|
||
data = tensor_to_cv2(data)
|
||
|
||
# 2. Ensure data is uint8 (OpenCV requirement for .png)
|
||
if data.dtype != np.uint8:
|
||
data = (data.clip(0, 1) * 255).astype(np.uint8) if data.max() <= 1.0 else data.astype(np.uint8)
|
||
|
||
# 3. Build full path
|
||
file_name = f"{name}_{suffix}.png"
|
||
save_path = os.path.join(self.output_dir, file_name)
|
||
|
||
# 4. Save and Check result
|
||
success = cv2.imwrite(save_path, data)
|
||
|
||
if success:
|
||
print(f"✓ Saved: {file_name}")
|
||
else:
|
||
print(f"× Failed to save: {file_name} (Check permissions or path)")
|
||
|
||
# =====================================================================================================================
|
||
# =====================================================================================================================
|
||
# =====================================================================================================================
|
||
|
||
class QFXTiledPreview:
|
||
def __init__(self):
|
||
self.output_dir = folder_paths.get_temp_directory()
|
||
self.type = "temp"
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE", {}),
|
||
"tile_x": ("INT", {"default": 3, "min": 1, "max": 10}),
|
||
"tile_y": ("INT", {"default": 3, "min": 1, "max": 10}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ()
|
||
RETURN_NAMES = ()
|
||
FUNCTION = "make_preview"
|
||
CATEGORY = "QFX"
|
||
OUTPUT_NODE = True
|
||
|
||
def make_preview(self, image, tile_x, tile_y):
|
||
|
||
batch_out = []
|
||
for img in image:
|
||
chw = _to_chw(img)
|
||
tile = _make_tile_preview(chw, tiles_x=tile_x, tiles_y=tile_y)
|
||
tile = _clamp01(tile)
|
||
batch_out.append(_to_hwc(tile))
|
||
|
||
result_tensor = torch.stack(batch_out, dim=0)
|
||
|
||
# --------------------------------------
|
||
|
||
# Save the first image in the batch for the UI preview
|
||
results = []
|
||
for tensor in result_tensor:
|
||
array = 255. * tensor.cpu().numpy()
|
||
img = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))
|
||
|
||
# Generate a unique filename
|
||
filename = f"QFX_Preview_{random.randint(0, 1000000)}.png"
|
||
img.save(os.path.join(self.output_dir, filename), compress_level=4)
|
||
results.append({
|
||
"filename": filename,
|
||
"subfolder": "",
|
||
"type": self.type
|
||
})
|
||
|
||
# Return both the tensor for downstream nodes AND the UI instruction
|
||
return {"ui": {"images": results}, "result": (result_tensor,)} |