__init__.py

import torch
import numpy as np
from PIL import Image
import io
import base64
from server import PromptServer
import threading
from aiohttp import web
import json
import torch.nn.functional as F
from threading import Event
image_cache = {}
event_dict = {}
CATEGORY_TYPE = "🎈LAOGOU/Relight"
class LG_Relight_Basic:

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "normals": ("IMAGE",),
                "x": ("FLOAT", { 
                    "default": 0.5, 
                    "min": 0.0, 
                    "max": 1.0, 
                    "step": 0.001 
                }),
                "y": ("FLOAT", { 
                    "default": 0.5, 
                    "min": 0.0, 
                    "max": 1.0, 
                    "step": 0.001 
                }),
                "z": ("FLOAT", { 
                    "default": 1.0, 
                    "min": -1.0, 
                    "max": 1.0, 
                    "step": 0.001 
                }),
                "brightness": ("FLOAT", { 
                    "default": 1.0, 
                    "min": 0.0, 
                    "max": 3.0, 
                    "step": 0.001 
                }),
                "shadow_range": ("FLOAT", { 
                    "default": 1.0, 
                    "min": 0.0, 
                    "max": 2.0, 
                    "step": 0.001 
                }),
                "shadow_strength": ("FLOAT", { 
                    "default": 1.0, 
                    "min": 0.0, 
                    "max": 2.0, 
                    "step": 0.001 
                }),
                "highlight_range": ("FLOAT", { 
                    "default": 1.0, 
                    "min": 0.0, 
                    "max": 2.0, 
                    "step": 0.001 
                }),
                "highlight_strength": ("FLOAT", { 
                    "default": 1.0, 
                    "min": 0.0, 
                    "max": 2.0, 
                    "step": 0.001 
                }),
                "highlight_color": ("STRING", {"default": "#FFFFFF"}),
                "shadow_color": ("STRING", {"default": "#000000"}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "relight"
    CATEGORY = CATEGORY_TYPE

    def relight(self, image, normals, x, y, z, brightness, 
              shadow_range, shadow_strength, highlight_range, highlight_strength,
              highlight_color, shadow_color):
        
        # 将十六进制颜色转换为RGB值
        def hex_to_rgb(hex_color):
            hex_color = hex_color.lstrip('#')
            return [int(hex_color[i:i+2], 16)/255.0 for i in (0, 2, 4)]
        
        highlight_color = hex_to_rgb(highlight_color)
        shadow_color = hex_to_rgb(shadow_color)
        
        norm = normals.detach().clone() * 2 - 1
        norm = F.interpolate(norm.movedim(-1,1), 
                        size=(image.shape[1], image.shape[2]), 
                        mode='bilinear').movedim(1,-1)

        # 将百分比坐标转换为光照方向
        # 确保光源位置与坐标一致
        light_x = -((x * 2) - 1)  # 0->-1, 1->1
        # 注意这里y要反转，因为图像坐标系y轴向下，而光照坐标系y轴向上
        light_y = -((y * 2) - 1)  # 0->1, 1->-1
        
        light = torch.tensor([light_x, light_y, z], device=image.device)
        light = F.normalize(light, dim=0)

        diffuse = norm[:,:,:,0] * light[0] + norm[:,:,:,1] * light[1] + norm[:,:,:,2] * light[2]
        diffuse = (diffuse + 1.0) * 0.5
        
        shadow_offset = shadow_strength - 1.0
        highlight_offset = highlight_strength - 1.0

        shadow_threshold = 1.0 - shadow_range
        highlight_threshold = 1.0 - highlight_range
        
        shadow_mask = torch.clamp((diffuse - shadow_threshold) / max(shadow_range, 1e-6), 0, 1)
        highlight_mask = torch.clamp((diffuse - highlight_threshold) / max(highlight_range, 1e-6), 0, 1)

        light_intensity = torch.ones_like(diffuse)

        if shadow_strength != 1.0:
            light_intensity = light_intensity * (
                shadow_mask + 
                (1.0 - shadow_mask) * (2.0 - shadow_strength)
            )

        if highlight_strength != 1.0:
            light_intensity = light_intensity + highlight_mask * highlight_offset

        color_effect = torch.ones_like(image[:,:,:,:3])
        if highlight_color != [1.0, 1.0, 1.0] or shadow_color != [0.0, 0.0, 0.0]:
            highlight_color = torch.tensor(highlight_color, device=image.device)
            shadow_color = torch.tensor(shadow_color, device=image.device)
            color_effect = (
                shadow_mask.unsqueeze(-1) * highlight_color +
                (1.0 - shadow_mask).unsqueeze(-1) * shadow_color
            )

        brightness_factor = brightness if brightness != 1.0 else 1.0

        relit = image.detach().clone()
        light_intensity = light_intensity.unsqueeze(-1).repeat(1,1,1,3)
        
        relit[:,:,:,:3] = torch.clip(
            relit[:,:,:,:3] * light_intensity * brightness_factor * color_effect,
            0, 1
        )
        
        return (relit,)
    


class LG_Relight:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "normals": ("IMAGE",),
            },
            "hidden": {
                "unique_id": "UNIQUE_ID",
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "relight"
    CATEGORY = CATEGORY_TYPE
    OUTPUT_NODE = True

    def encode_image_to_base64(self, image, is_mask=False):
        image = (image * 255).clip(0, 255).astype(np.uint8)
        
        if is_mask:
            if len(image.shape) == 3:
                image = image[0]
            image = np.stack([image] * 3, axis=-1)
        else:
            if len(image.shape) == 4:
                image = image[0]
        
        image = Image.fromarray(image)
        buffered = io.BytesIO()
        image.save(buffered, format="PNG")
        return base64.b64encode(buffered.getvalue()).decode()

    def relight(self, image, normals, unique_id):
        try:
            event = threading.Event()
            event_dict[unique_id] = event
            image_b64 = self.encode_image_to_base64(image.cpu().numpy())
            normals_b64 = self.encode_image_to_base64(normals.cpu().numpy())
            send_data = {
                "node_id": unique_id,
                "image": f"data:image/png;base64,{image_b64}",
                "normals": f"data:image/png;base64,{normals_b64}"
            }
            PromptServer.instance.send_sync("lg_relight_init", send_data)
            event.wait()
            del event_dict[unique_id]
            
            if unique_id in image_cache:
                img_data = base64.b64decode(image_cache[unique_id].split(",")[1])
                img = Image.open(io.BytesIO(img_data))
                img_np = np.array(img).astype(np.float32) / 255.0
                
                if len(img_np.shape) == 2:
                    img_np = np.stack([img_np] * 3, axis=-1)
                elif len(img_np.shape) == 3 and img_np.shape[-1] == 4:
                    img_np = img_np[..., :3]
                
                result = torch.from_numpy(img_np).unsqueeze(0)
                
                del image_cache[unique_id]
                return (result,)
            else:
                return (image,)
            
        except Exception as e:
            print(f"[ERROR] An error occurred during relight processing: {str(e)}")
            import traceback
            print(traceback.format_exc())
            return (image,)
        finally:
            print(f"[DEBUG] 清理资源Cleaning up resources: node_id={unique_id}")
            if unique_id in event_dict:
                del event_dict[unique_id]
            if unique_id in image_cache:
                del image_cache[unique_id]

@PromptServer.instance.routes.post("/lg_relight/update_image")
async def update_image_v3(request):
    try:
        data = await request.json()
        node_id = data.get("node_id")
        image_data = data.get("image")
        
        if node_id and image_data:
            image_cache[node_id] = image_data
            if node_id in event_dict:
                event_dict[node_id].set()
            return web.Response(text=json.dumps({"status": "success"}))
        else:
            return web.Response(status=400, text=json.dumps({"error": "无效数据Invalid data"}))
    except Exception as e:
        return web.Response(status=500, text=json.dumps({"error": str(e)}))

@PromptServer.instance.routes.post("/lg_relight/cancel")
async def cancel_v3(request):
    try:
        data = await request.json()
        node_id = data.get("node_id")
        
        if node_id:
            if node_id in event_dict:
                event_dict[node_id].set()
            return web.Response(text=json.dumps({"status": "success"}))
        else:
            return web.Response(status=400, text=json.dumps({"error": "无效节点ID"}))
    except Exception as e:
        return web.Response(status=500, text=json.dumps({"error": str(e)}))


lg_relight_dict = {}

class LG_Relight_Ultra:
    _last_results = {}
    
    def __init__(self):
        self.node_id = None
    
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "bg_img": ("IMAGE",),
                "bg_depth_map": ("IMAGE",),
                "bg_normal_map": ("IMAGE",),
                "wait_timeout": ("INT", {
                    "default": 120,
                    "min": 5,
                    "max": 300,
                    "step": 1,
                    "tooltip": "等待前端响应的最大时间(秒)\nMaximum time to wait for frontend response (seconds)"
                }),
            },
            "optional": {
                "mask": ("MASK",),
                "skip_dialog": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "开启后将不再显示光照编辑窗口，直接使用之前保存的光照设置（如果没有则使用默认设置）\nEnable to skip the lighting editor dialog and directly use previously saved lighting settings (or default settings if none exist)"
                }),
            },
            "hidden": {"unique_id": "UNIQUE_ID"}
        }
    
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "relight_image"
    CATEGORY = CATEGORY_TYPE

    def relight_image(self, bg_img, bg_depth_map, bg_normal_map, wait_timeout, unique_id, mask=None, skip_dialog=False):
        try:
            self.node_id = str(unique_id)
            event = Event()
            lg_relight_dict[self.node_id] = event
            
            bg_pil = Image.fromarray((bg_img[0] * 255).byte().cpu().numpy())
            depth_pil = Image.fromarray((bg_depth_map[0] * 255).byte().cpu().numpy())
            normal_pil = Image.fromarray((bg_normal_map[0] * 255).byte().cpu().numpy())

            bg_buffer = io.BytesIO()
            depth_buffer = io.BytesIO()
            normal_buffer = io.BytesIO()
            
            bg_pil.save(bg_buffer, format="PNG")
            depth_pil.save(depth_buffer, format="PNG")
            normal_pil.save(normal_buffer, format="PNG")
            
            data = {
                "node_id": self.node_id,
                "bg_image": base64.b64encode(bg_buffer.getvalue()).decode('utf-8'),
                "bg_depth_map": base64.b64encode(depth_buffer.getvalue()).decode('utf-8'),
                "bg_normal_map": base64.b64encode(normal_buffer.getvalue()).decode('utf-8'),
                "has_mask": mask is not None,
                "skip_dialog": skip_dialog
            }
            
            if mask is not None:
                try:
                    mask_np = mask
                    if isinstance(mask, torch.Tensor):
                        mask_np = (mask * 255).byte().cpu().numpy()
                    
                    if len(mask_np.shape) == 3 and mask_np.shape[0] == 1:
                        mask_np = mask_np[0]
                    elif len(mask_np.shape) == 4 and mask_np.shape[0] == 1:
                        mask_np = mask_np[0]
                    
                    if mask_np.dtype != np.uint8:
                        mask_np = (mask_np * 255).astype(np.uint8)
                    
                    mask_pil = Image.fromarray(mask_np)
                    mask_buffer = io.BytesIO()
                    mask_pil.save(mask_buffer, format="PNG")
                    data["mask"] = base64.b64encode(mask_buffer.getvalue()).decode('utf-8')
                    
                except Exception:
                    data["has_mask"] = False

            PromptServer.instance.send_sync("relight_image", data)
            
            wait_result = event.wait(timeout=wait_timeout)
            if not wait_result:
                return (bg_img,)
            
            if self.node_id in self._last_results:
                result_image = self._last_results[self.node_id]
                try:
                    img = Image.open(io.BytesIO(result_image))
                    img_array = np.array(img)
                    
                    if len(img_array.shape) == 2:
                        img_array = np.stack([img_array] * 3, axis=-1)
                    elif img_array.shape[-1] == 4:
                        img_array = img_array[..., :3]
                    
                    img_tensor = torch.from_numpy(img_array).float() / 255.0
                    img_tensor = img_tensor.unsqueeze(0)
                    return (img_tensor,)
                except Exception:
                    return (bg_img,)
            
            return (bg_img,)
            
        finally:
            if self.node_id in lg_relight_dict:
                del lg_relight_dict[self.node_id]

@PromptServer.instance.routes.post("/lg_relight/upload_result")
async def upload_result(request):
    try:
        data = await request.post()
        node_id = str(data['node_id'])
        result_image = data['result_image'].file.read()
        
        LG_Relight_Ultra._last_results[node_id] = result_image
        
        if node_id in lg_relight_dict:
            lg_relight_dict[node_id].set()
            return web.json_response({"success": True})
        else:
            return web.json_response({"error": "Node not found"}, status=404)
            
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)

@PromptServer.instance.routes.post("/lg_relight_ultra/cancel")
async def cancel_relight(request):
    try:
        data = await request.json()
        node_id = str(data.get("node_id"))
        
        if node_id in lg_relight_dict:
            lg_relight_dict[node_id].set()
            return web.json_response({"success": True})
        else:
            return web.json_response({"success": True})
        
    except Exception as e:
        return web.json_response({"success": False, "error": str(e)})

WEB_DIRECTORY = "web"

NODE_CLASS_MAPPINGS = {
    "LG_Relight_Basic": LG_Relight_Basic,
    "LG_Relight": LG_Relight,
    "LG_Relight_Ultra": LG_Relight_Ultra
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LG_Relight_Basic": "🎈LG Relight Basic",
    "LG_Relight": "🎈LG Relight",
    "LG_Relight_Ultra": "🎈LG Relight Ultra"
}