Fixed Reactor to Pytorch2.6.0 on ComfyUI

import numpy as np
import cv2
from PIL import Image
import torch
import torch.serialization
import scripts.r_masking.core as core
from reactor_utils import tensor_to_pil
import logging

# 警告メッセージを抑制するために一時的にロギングレベルを変更
original_level = logging.getLogger().level
logging.getLogger().setLevel(logging.ERROR)

try:
    from ultralytics import YOLO
    # ultralytics関連のクラス
    from ultralytics.nn.tasks import DetectionModel
    from ultralytics.nn.modules.conv import Conv
    
    # 基本的なブロックとヘッド - 必須のクラス
    basic_classes = []
    try:
        from ultralytics.nn.modules.block import C2f, C3, C2, SPPF, Bottleneck, BottleneckCSP
        from ultralytics.nn.modules.head import Detect, Segment, Pose, OBB, Classify
        basic_classes = [C2f, C3, C2, SPPF, Bottleneck, BottleneckCSP, Detect, Segment, Pose, OBB, Classify]
    except ImportError:
        pass
    
    # torch関連の標準クラス
    from torch.nn.modules.container import Sequential, ModuleList
    from torch.nn.modules.conv import Conv2d, ConvTranspose2d
    from torch.nn.modules.batchnorm import BatchNorm2d
    from torch.nn.modules.activation import SiLU, ReLU, LeakyReLU, Hardswish
    from torch.nn.modules.linear import Linear, Identity
    from torch.nn.modules.dropout import Dropout
    from torch.nn.modules.pooling import MaxPool2d, AvgPool2d
    from torch.nn.modules.padding import ZeroPad2d, ReflectionPad2d
    from torch.nn.modules.upsampling import Upsample
    from torch.nn.modules.normalization import GroupNorm, LayerNorm
    
    # 追加のultralytics特有のクラス - 静かに失敗させる
    additional_classes = []
    try:
        from ultralytics.nn.modules.block import DFL, Proto, RepC3, C3TR, C3Ghost, GhostConv, HGBlock, SPP
        additional_classes.extend([DFL, Proto, RepC3, C3TR, C3Ghost, GhostConv, HGBlock, SPP])
    except ImportError:
        pass
        
    try:
        from ultralytics.nn.modules.block import RepConv, Focus, ChannelAttention, SpatialAttention, CBAM
        additional_classes.extend([RepConv, Focus, ChannelAttention, SpatialAttention, CBAM])
    except ImportError:
        pass
        
    try:
        from ultralytics.nn.modules.block import Concat, TransformerBlock, C3STR, C2f_DCN, Space2Depth
        additional_classes.extend([Concat, TransformerBlock, C3STR, C2f_DCN, Space2Depth])
    except ImportError:
        pass
    
    # 必要なクラスをセーフリストに追加
    safe_classes = [
        # ultralytics関連
        DetectionModel, Conv,
        
        # torch関連
        Sequential, ModuleList, Conv2d, ConvTranspose2d, BatchNorm2d,
        SiLU, ReLU, LeakyReLU, Hardswish, Linear, Identity, Dropout,
        MaxPool2d, AvgPool2d, ZeroPad2d, ReflectionPad2d, Upsample,
        GroupNorm, LayerNorm
    ]
    
    # 基本クラスと追加クラスを追加
    safe_classes.extend(basic_classes)
    safe_classes.extend(additional_classes)
    
    # セーフリストに追加
    torch.serialization.add_safe_globals(safe_classes)
    
except Exception as e:
    # エラーは記録するが、警告は表示しない
    logging.error(f"YOLOモデルの初期化中にエラーが発生しました: {e}")

# ロギングレベルを元に戻す
logging.getLogger().setLevel(original_level)

# 代替方法: weights_only=Falseを使用する関数
def load_yolo(model_path: str):
    try:
        # 通常の方法でロードを試みる
        return YOLO(model_path)
    except Exception as e:
        # エラーをログに記録するが、コンソールには表示しない
        logging.error(f"通常のロード方法でエラーが発生しました: {e}")
        
        # 代替方法: weights_only=Falseを使用
        original_torch_load = torch.load
        try:
            # 一時的にtorch.loadを上書き
            torch.load = lambda *args, **kwargs: original_torch_load(*args, weights_only=False, **kwargs)
            model = YOLO(model_path)
            return model
        except ModuleNotFoundError:
            # https://github.com/ultralytics/ultralytics/issues/3856
            YOLO("yolov8n.pt")
            return YOLO(model_path)
        finally:
            # 元のtorch.loadに戻す
            torch.load = original_torch_load

# ... 残りのコードは変更なし ...

def inference_bbox(
    model,
    image: Image.Image,
    confidence: float = 0.3,
    device: str = "",
):
    pred = model(image, conf=confidence, device=device)

    bboxes = pred[0].boxes.xyxy.cpu().numpy()
    cv2_image = np.array(image)
    if len(cv2_image.shape) == 3:
        cv2_image = cv2_image[:, :, ::-1].copy()  # Convert RGB to BGR for cv2 processing
    else:
        # Handle the grayscale image here
        # For example, you might want to convert it to a 3-channel grayscale image for consistency:
        cv2_image = cv2.cvtColor(cv2_image, cv2.COLOR_GRAY2BGR)
    cv2_gray = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2GRAY)

    segms = []
    for x0, y0, x1, y1 in bboxes:
        cv2_mask = np.zeros(cv2_gray.shape, np.uint8)
        cv2.rectangle(cv2_mask, (int(x0), int(y0)), (int(x1), int(y1)), 255, -1)
        cv2_mask_bool = cv2_mask.astype(bool)
        segms.append(cv2_mask_bool)

    n, m = bboxes.shape
    if n == 0:
        return [[], [], [], []]

    results = [[], [], [], []]
    for i in range(len(bboxes)):
        results[0].append(pred[0].names[int(pred[0].boxes[i].cls.item())])
        results[1].append(bboxes[i])
        results[2].append(segms[i])
        results[3].append(pred[0].boxes[i].conf.cpu().numpy())

    return results


class UltraBBoxDetector:
    bbox_model = None

    def __init__(self, bbox_model):
        self.bbox_model = bbox_model

    def detect(self, image, threshold, dilation, crop_factor, drop_size=1, detailer_hook=None):
        drop_size = max(drop_size, 1)
        detected_results = inference_bbox(self.bbox_model, tensor_to_pil(image), threshold)
        segmasks = core.create_segmasks(detected_results)

        if dilation > 0:
            segmasks = core.dilate_masks(segmasks, dilation)

        items = []
        h = image.shape[1]
        w = image.shape[2]

        for x, label in zip(segmasks, detected_results[0]):
            item_bbox = x[0]
            item_mask = x[1]

            y1, x1, y2, x2 = item_bbox

            if x2 - x1 > drop_size and y2 - y1 > drop_size:  # minimum dimension must be (2,2) to avoid squeeze issue
                crop_region = core.make_crop_region(w, h, item_bbox, crop_factor)

                if detailer_hook is not None:
                    crop_region = detailer_hook.post_crop_region(w, h, item_bbox, crop_region)

                cropped_image = core.crop_image(image, crop_region)
                cropped_mask = core.crop_ndarray2(item_mask, crop_region)
                confidence = x[2]
                # bbox_size = (item_bbox[2]-item_bbox[0],item_bbox[3]-item_bbox[1]) # (w,h)

                item = core.SEG(cropped_image, cropped_mask, confidence, crop_region, item_bbox, label, None)

                items.append(item)

        shape = image.shape[1], image.shape[2]
        segs = shape, items

        if detailer_hook is not None and hasattr(detailer_hook, "post_detection"):
            segs = detailer_hook.post_detection(segs)

        return segs

    def detect_combined(self, image, threshold, dilation):
        detected_results = inference_bbox(self.bbox_model, core.tensor2pil(image), threshold)
        segmasks = core.create_segmasks(detected_results)
        if dilation > 0:
            segmasks = core.dilate_masks(segmasks, dilation)

        return core.combine_masks(segmasks)

    def setAux(self, x):
        pass