import math
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

import comfy.model_management
import comfy.ops
import comfy.text_encoders.qwen_vl
from comfy import sd1_clip
from comfy.ldm.rt_detr.rtdetr_v4 import SelfAttention, _deformable_attn_v2
from comfy.text_encoders.llama import Llama2_, Qwen3VL_4BConfig
from comfy.text_encoders.qwen3vl import QWEN3VL_VISION, QWEN3VL_VISION_COMMON, Qwen3VL, Qwen3VLClipModel, Qwen3VLTokenizer, Qwen3VLVisionModel

# Tokens appended to the Qwen vocabulary, in id order starting at 151669.
ADDED_TOKENS = ["<{}>".format(i) for i in range(1001)] + ["<ref>", "</ref>", "<box>", "</box>", "<point>", "</point>", "<polygon>", "</polygon>", "<OBJ>"]
OBJ_TOKEN = 152678
SYSTEM_PROMPT = "You are a helpful assistant."
MIN_PIXELS = 16 * 32 * 32
MAX_PIXELS = 2025 * 32 * 32
FEATURE_LAYERS = (5, 11, 17, 23)
CNN_SIZE = 1024
# Peak activations of one mask query at the 1024x1024 decoder canvas.
DECODER_MEMORY = 2 * 1024 ** 3


@dataclass
class WeVisOneConfig(Qwen3VL_4BConfig):
    vocab_size: int = 152680
    lm_head: bool = True


# Pixel decoder (mask_decoder.pt), hidden 192, mask features 64, 4 ViT levels, 6 query layers.
def group_norm(channels, groups, device, dtype, ops):
    groups = min(groups, channels)
    while channels % groups:
        groups -= 1
    return ops.GroupNorm(groups, channels, device=device, dtype=dtype)


def fourier(values, num_freqs=32):
    freqs = 2.0 ** torch.arange(num_freqs, device=values.device, dtype=torch.float32)
    phase = 2.0 * math.pi * values[..., None] * freqs
    return torch.cat([phase.sin(), phase.cos()], dim=-1)


def sine_position(h, w, dim, device):
    half = dim // 2
    y = (torch.arange(h, device=device, dtype=torch.float32) + 0.5) / h * (2.0 * math.pi)
    x = (torch.arange(w, device=device, dtype=torch.float32) + 0.5) / w * (2.0 * math.pi)
    yy, xx = torch.meshgrid(y, x, indexing="ij")
    dim_t = 10000.0 ** (2.0 * (torch.arange(half, device=device, dtype=torch.float32) // 2) / half)
    pos_x = xx[..., None] / dim_t
    pos_y = yy[..., None] / dim_t
    pos_x = torch.stack([pos_x[..., 0::2].sin(), pos_x[..., 1::2].cos()], dim=-1).flatten(-2)
    pos_y = torch.stack([pos_y[..., 0::2].sin(), pos_y[..., 1::2].cos()], dim=-1).flatten(-2)
    return torch.cat([pos_y, pos_x], dim=-1).reshape(h * w, dim)


def resize(x, size):
    if tuple(x.shape[-2:]) == tuple(size):
        return x
    return F.interpolate(x, size=size, mode="bilinear", align_corners=False)


class ConvGN(nn.Sequential):
    def __init__(self, in_ch, out_ch, kernel, groups, stride=1, act=True, device=None, dtype=None, ops=None):
        layers = [ops.Conv2d(in_ch, out_ch, kernel, stride=stride, padding=kernel // 2, bias=False, device=device, dtype=dtype), group_norm(out_ch, groups, device, dtype, ops)]
        if act:
            layers.append(nn.GELU())
        super().__init__(*layers)


class BoxFourierEncoder(nn.Module):
    def __init__(self, hidden, device=None, dtype=None, ops=None):
        super().__init__()
        self.mlp = nn.Sequential(ops.Linear(4 * 64, hidden, device=device, dtype=dtype), nn.GELU(), ops.Linear(hidden, hidden, device=device, dtype=dtype))

    def forward(self, boxes):
        x1, y1, x2, y2 = boxes.unbind(-1)
        w = 0.2 * (x2 - x1).clamp_min(1e-6).log() + 0.5
        h = 0.2 * (y2 - y1).clamp_min(1e-6).log() + 0.5
        return self.mlp(torch.cat([fourier(x1), fourier(y1), fourier(w), fourier(h)], dim=-1))


class MaskQueryBuilder(nn.Module):
    def __init__(self, llm_dim, hidden, device=None, dtype=None, ops=None):
        super().__init__()
        self.content_proj = ops.Linear(llm_dim, hidden, device=device, dtype=dtype)
        self.box_position = BoxFourierEncoder(hidden, device=device, dtype=dtype, ops=ops)
        self.content_norm = ops.LayerNorm(hidden, device=device, dtype=dtype)
        self.position_norm = ops.LayerNorm(hidden, device=device, dtype=dtype)

    def forward(self, t_seg, boxes):
        return self.content_norm(self.content_proj(t_seg)), self.position_norm(self.box_position(boxes))


class SpatialConvEncoderLayer(nn.Module):
    def __init__(self, d=192, ffn=768, device=None, dtype=None, ops=None):
        super().__init__()
        self.depthwise = ops.Conv2d(d, d, 3, padding=1, groups=d, bias=False, device=device, dtype=dtype)
        self.pointwise = ops.Conv2d(d, d, 1, bias=False, device=device, dtype=dtype)
        self.norm1 = ops.LayerNorm(d, device=device, dtype=dtype)
        self.linear1 = ops.Linear(d, ffn, device=device, dtype=dtype)
        self.linear2 = ops.Linear(ffn, d, device=device, dtype=dtype)
        self.norm2 = ops.LayerNorm(d, device=device, dtype=dtype)

    def forward(self, src, pos, h, w):
        levels = (src + pos).transpose(1, 2).unflatten(2, (-1, h, w)).transpose(1, 2).flatten(0, 1)
        mixed = self.pointwise(F.gelu(self.depthwise(levels)))
        mixed = mixed.unflatten(0, (src.shape[0], -1)).flatten(3).permute(0, 1, 3, 2).flatten(1, 2)
        src = self.norm1(src + mixed)
        return self.norm2(src + self.linear2(F.gelu(self.linear1(src))))


class VisualEncoder(nn.Module):
    def __init__(self, in_dim, d=192, levels=4, device=None, dtype=None, ops=None):
        super().__init__()
        self.input_proj = nn.ModuleList([ConvGN(in_dim, d, 1, 32, act=False, device=device, dtype=dtype, ops=ops) for _ in range(levels)])
        self.level_embed = nn.Parameter(torch.empty(levels, d, device=device, dtype=dtype))
        self.encoder = nn.Module()
        self.encoder.layers = nn.ModuleList([SpatialConvEncoderLayer(d, device=device, dtype=dtype, ops=ops) for _ in range(3)])

    def forward(self, features):
        # All levels share the pre-merge ViT grid.
        h, w = features[0].shape[-2:]
        projected = [proj(f).flatten(2).transpose(1, 2) for proj, f in zip(self.input_proj, features)]
        src = torch.cat(projected, dim=1)
        pos = sine_position(h, w, src.shape[-1], src.device).to(src)
        level_embed = comfy.ops.cast_to_input(self.level_embed, src)
        pos = torch.cat([pos + level_embed[i] for i in range(len(projected))], dim=0).unsqueeze(0)
        for layer in self.encoder.layers:
            src = layer(src, pos, h, w)
        maps = [chunk.transpose(1, 2).unflatten(2, (h, w)) for chunk in src.chunk(len(projected), dim=1)]
        return maps, src, h, w


class MSDeformableAttention(nn.Module):
    def __init__(self, d=192, heads=8, levels=4, points=4, device=None, dtype=None, ops=None):
        super().__init__()
        self.heads, self.levels, self.points = heads, levels, points
        self.sampling_offsets = ops.Linear(d, heads * levels * points * 2, device=device, dtype=dtype)
        self.attention_weights = ops.Linear(d, heads * levels * points, device=device, dtype=dtype)
        self.value_proj = ops.Linear(d, d, device=device, dtype=dtype)
        self.output_proj = ops.Linear(d, d, device=device, dtype=dtype)

    def forward(self, query, references, memory, h, w):
        b, q, _ = query.shape
        value = self.value_proj(memory).unflatten(2, (self.heads, -1)).permute(0, 2, 3, 1).flatten(0, 1)
        value = list(value.unflatten(2, (self.levels, h, w)).unbind(2))
        offsets = self.sampling_offsets(query).view(b, q, self.heads, self.levels * self.points, 2)
        weights = self.attention_weights(query).view(b, q, self.heads, -1).softmax(-1)
        locations = references[:, :, None, None, :2] + offsets / self.points * references[:, :, None, None, 2:] * 0.5
        out = _deformable_attn_v2(value, [(h, w)] * self.levels, locations, weights, [self.points] * self.levels)
        return self.output_proj(out)


class QueryDecoderLayer(nn.Module):
    def __init__(self, d=192, heads=8, ffn=768, device=None, dtype=None, ops=None):
        super().__init__()
        self.self_attn = SelfAttention(d, heads, device=device, dtype=dtype, operations=ops)
        self.cross_attn = MSDeformableAttention(d, heads, device=device, dtype=dtype, ops=ops)
        self.norm1 = ops.LayerNorm(d, device=device, dtype=dtype)
        self.norm2 = ops.LayerNorm(d, device=device, dtype=dtype)
        self.norm3 = ops.LayerNorm(d, device=device, dtype=dtype)
        self.linear1 = ops.Linear(d, ffn, device=device, dtype=dtype)
        self.linear2 = ops.Linear(ffn, d, device=device, dtype=dtype)

    def forward(self, target, query_pos, references, memory, h, w):
        qk = target + query_pos
        target = self.norm1(target + self.self_attn(qk, qk, target))
        target = self.norm2(target + self.cross_attn(target + query_pos, references, memory, h, w))
        return self.norm3(target + self.linear2(F.gelu(self.linear1(target))))


class PixelShuffleStage(nn.Module):
    def __init__(self, in_dim, out_dim, device=None, dtype=None, ops=None):
        super().__init__()
        self.conv = ops.Conv2d(in_dim, out_dim * 4, 3, padding=1, device=device, dtype=dtype)

    def forward(self, x):
        return F.gelu(F.pixel_shuffle(self.conv(x), 2))


class ResidualPixelShuffleStage(PixelShuffleStage):
    def forward(self, x):
        return F.interpolate(x, scale_factor=2.0, mode="bilinear", align_corners=False) + F.pixel_shuffle(self.conv(x), 2)


class SameScaleFPN(nn.Module):
    def __init__(self, d, levels, device=None, dtype=None, ops=None):
        super().__init__()
        self.lateral = nn.ModuleList([ops.Conv2d(d, d, 1, bias=False, device=device, dtype=dtype) for _ in range(levels)])
        self.output = nn.ModuleList([ConvGN(d, d, 3, 32, device=device, dtype=dtype, ops=ops) for _ in range(levels - 1)])

    def forward(self, features):
        fused = self.lateral[-1](features[-1])
        for out, level in zip(self.output, range(len(features) - 2, -1, -1)):
            fused = out(fused + self.lateral[level](features[level]))
        return fused


class PixelFusion(nn.Module):
    def __init__(self, d=192, mask_dim=64, device=None, dtype=None, ops=None):
        super().__init__()
        self.fpn = SameScaleFPN(d, 4, device=device, dtype=dtype, ops=ops)
        self.upsample = nn.Sequential(PixelShuffleStage(d, d, device=device, dtype=dtype, ops=ops), PixelShuffleStage(d, mask_dim, device=device, dtype=dtype, ops=ops))
        self.s2_upsample = ResidualPixelShuffleStage(mask_dim, mask_dim, device=device, dtype=dtype, ops=ops)
        self.cnn = nn.Module()
        self.cnn.net = nn.Sequential(*ConvGN(3, 128, 3, 32, stride=2, device=device, dtype=dtype, ops=ops), *ConvGN(128, mask_dim, 3, 32, device=device, dtype=dtype, ops=ops))
        self.fpn_proj = ops.Conv2d(mask_dim, mask_dim, 1, bias=False, device=device, dtype=dtype)
        self.cnn_proj = ops.Conv2d(mask_dim, mask_dim, 1, bias=False, device=device, dtype=dtype)
        self.smooth = nn.Sequential(group_norm(mask_dim, 32, device, dtype, ops), nn.GELU(), ops.Conv2d(mask_dim, mask_dim, 3, padding=1, bias=False, device=device, dtype=dtype))

    def forward(self, image, levels):
        fpn = self.upsample(self.fpn(levels))
        cnn = self.cnn.net(image)
        target = (cnn.shape[-2] // 2, cnn.shape[-1] // 2)
        pool = (min(fpn.shape[-2], target[0]), min(fpn.shape[-1], target[1]))
        if pool != tuple(fpn.shape[-2:]):
            fpn = F.adaptive_avg_pool2d(fpn, pool)
        fpn = resize(fpn, target)
        return self.s2_upsample(self.smooth(self.fpn_proj(fpn))) + self.cnn_proj(cnn)


class AlphaResidualBlock(nn.Module):
    def __init__(self, ch, device=None, dtype=None, ops=None):
        super().__init__()
        self.net = nn.Sequential(group_norm(ch, 32, device, dtype, ops), nn.GELU(), ops.Conv2d(ch, ch, 3, padding=1, bias=False, device=device, dtype=dtype),
                                 group_norm(ch, 32, device, dtype, ops), nn.GELU(), ops.Conv2d(ch, ch, 3, padding=1, bias=False, device=device, dtype=dtype))

    def forward(self, x):
        return x + self.net(x)


class AlphaCascadeHead(nn.Module):
    # Refines the S/2 coarse logits at S/4, S/2 and S/1 with RGB, semantic features and the previous alpha.
    def __init__(self, d=192, mask_dim=64, ch=48, device=None, dtype=None, ops=None):
        super().__init__()
        project = lambda i: ConvGN(i, ch, 3, 32, device=device, dtype=dtype, ops=ops)
        stack = lambda: nn.Sequential(*[AlphaResidualBlock(ch, device=device, dtype=dtype, ops=ops) for _ in range(2)])
        self.semantic = project(mask_dim)
        self.rgb_s4 = project(3)
        self.rgb_s2 = project(3)
        self.rgb_s1 = project(3)
        self.coarse_s4 = project(1)
        self.coarse_s2 = project(1)
        self.previous_s2 = project(1)
        self.previous_s1 = project(1)
        self.fuse_s4 = ops.Conv2d(3 * ch, ch, 1, bias=False, device=device, dtype=dtype)
        self.fuse_s2 = ops.Conv2d(5 * ch, ch, 1, bias=False, device=device, dtype=dtype)
        self.fuse_s1 = ops.Conv2d(4 * ch, ch, 1, bias=False, device=device, dtype=dtype)
        self.blocks_s4 = stack()
        self.blocks_s2 = stack()
        self.blocks_s1 = stack()
        self.up_s4_to_s2 = ops.ConvTranspose2d(ch, ch, 4, stride=2, padding=1, bias=False, device=device, dtype=dtype)
        self.up_s2_to_s1 = ops.ConvTranspose2d(ch, ch, 4, stride=2, padding=1, bias=False, device=device, dtype=dtype)
        self.query_s4 = ops.Linear(d, 2 * ch, device=device, dtype=dtype)
        self.query_s2 = ops.Linear(d, 2 * ch, device=device, dtype=dtype)
        self.query_s1 = ops.Linear(d, 2 * ch, device=device, dtype=dtype)
        self.output_s4 = ops.Conv2d(ch, 1, 1, device=device, dtype=dtype)
        self.output_s2 = ops.Conv2d(ch, 1, 1, device=device, dtype=dtype)
        self.output_s1 = ops.Conv2d(ch, 1, 1, device=device, dtype=dtype)

    @staticmethod
    def condition(x, queries, affine):
        scale, bias = affine(queries)[:, :, None, None].chunk(2, dim=1)
        return x * (1.0 + 0.1 * scale.tanh()) + 0.1 * bias

    def forward(self, image, pixel_feature, queries, coarse):
        n = queries.shape[0]
        s2 = tuple(coarse.shape[-2:])
        s4 = (max(1, s2[0] // 2), max(1, s2[1] // 2))
        s1 = tuple(image.shape[-2:])
        coarse = coarse[:, None]

        semantic_s2 = self.semantic(pixel_feature)
        coarse_s4 = resize(coarse, s4)
        x = self.fuse_s4(torch.cat([resize(semantic_s2, s4).expand(n, -1, -1, -1), self.rgb_s4(resize(image, s4)).expand(n, -1, -1, -1), self.coarse_s4(coarse_s4)], dim=1))
        feature_s4 = self.blocks_s4(self.condition(x, queries, self.query_s4))
        logits_s4 = coarse_s4 + self.output_s4(feature_s4)

        x = self.fuse_s2(torch.cat([self.up_s4_to_s2(feature_s4), semantic_s2.expand(n, -1, -1, -1), self.rgb_s2(resize(image, s2)).expand(n, -1, -1, -1),
                                    self.coarse_s2(coarse), self.previous_s2(resize(logits_s4, s2).sigmoid())], dim=1))
        feature_s2 = self.blocks_s2(self.condition(x, queries, self.query_s2))
        logits_s2 = coarse + self.output_s2(feature_s2)

        previous_s1 = resize(logits_s2, s1)
        x = self.fuse_s1(torch.cat([self.up_s2_to_s1(feature_s2), resize(semantic_s2, s1).expand(n, -1, -1, -1), self.rgb_s1(image).expand(n, -1, -1, -1),
                                    self.previous_s1(previous_s1.sigmoid())], dim=1))
        feature_s1 = self.blocks_s1(self.condition(x, queries, self.query_s1))
        return (previous_s1 + self.output_s1(feature_s1))[:, 0]


class WeVisOneMaskDecoder(nn.Module):
    def __init__(self, visual_dim=1024, llm_dim=2560, d=192, mask_dim=64, device=None, dtype=None, ops=None):
        super().__init__()
        self.visual_encoder = VisualEncoder(visual_dim, d, device=device, dtype=dtype, ops=ops)
        self.query_builder = MaskQueryBuilder(llm_dim, d, device=device, dtype=dtype, ops=ops)
        self.transformer = nn.Module()
        self.transformer.layers = nn.ModuleList([QueryDecoderLayer(d, device=device, dtype=dtype, ops=ops) for _ in range(6)])
        self.pixel_fusion = PixelFusion(d, mask_dim, device=device, dtype=dtype, ops=ops)
        self.mask_predictor = nn.Module()
        self.mask_predictor.kernel = nn.Sequential(ops.Linear(d, d, device=device, dtype=dtype), nn.GELU(), ops.Linear(d, mask_dim + 1, device=device, dtype=dtype))
        self.highres_refine = AlphaCascadeHead(d, mask_dim, device=device, dtype=dtype, ops=ops)

    def forward(self, image, visual_features, t_seg, boxes):
        # image [1, 3, 1024, 1024] in 0..1, visual_features 4x [1, C, gh, gw], t_seg [N, llm_dim], boxes [N, 4] normalized xyxy
        levels, memory, h, w = self.visual_encoder(visual_features)
        target, query_pos = self.query_builder(t_seg, boxes)
        x1, y1, x2, y2 = boxes.clamp(0.0, 1.0).unbind(-1)
        references = torch.stack([(x1 + x2) * 0.5, (y1 + y2) * 0.5, (x2 - x1).clamp_min(1e-4), (y2 - y1).clamp_min(1e-4)], dim=-1)[None]
        target, query_pos = target[None], query_pos[None]
        for layer in self.transformer.layers:
            target = layer(target, query_pos, references, memory, h, w)
        queries = target[0]
        pixel_feature = self.pixel_fusion(image, levels)
        params = self.mask_predictor.kernel(queries)
        coarse = torch.einsum("qc,chw->qhw", params[:, :-1], pixel_feature[0]) + params[:, -1, None, None]
        return self.highres_refine(image, pixel_feature, queries, coarse)


class WeVisOneQwen3VL(Qwen3VL):
    model_type = "qwen3vl_4b"

    def __init__(self, config_dict, dtype, device, operations):
        super(Qwen3VL, self).__init__()
        config = WeVisOneConfig(**config_dict)
        self.num_layers = config.num_hidden_layers
        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
        vision_config = {**QWEN3VL_VISION_COMMON, **QWEN3VL_VISION[self.model_type], "out_hidden_size": config.hidden_size}
        self.visual = Qwen3VLVisionModel(vision_config, device=device, dtype=dtype, ops=operations)
        # The released decoder weights are fp32 and the reference runs it in fp32. It is never quantized, so it
        # skips the mixed precision ops of quantized checkpoints, which would store its linears in the compute dtype.
        self.mask_decoder = WeVisOneMaskDecoder(device=device, dtype=torch.float32, ops=comfy.ops.manual_cast)
        self.dtype = dtype

    def preprocess_embed(self, embed, device):
        if embed["type"] != "image":
            return None, None
        image, grid = comfy.text_encoders.qwen_vl.process_qwen2vl_images(embed["data"], min_pixels=MIN_PIXELS, max_pixels=MAX_PIXELS, patch_size=16, image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5], interpolation="bicubic")
        # The pixel decoder reads layers 5/11/17/23 at the inputs of the DeepStack mergers and the final merger.
        premerge = []
        hook = lambda module, args: premerge.append(args[0])
        handles = [m.register_forward_pre_hook(hook) for m in list(self.visual.deepstack_merger_list) + [self.visual.merger]]
        try:
            merged, deepstack = self.visual(image.to(device, dtype=torch.float32), grid)
        finally:
            for h in handles:
                h.remove()
        return merged, {"grid": grid, "deepstack": deepstack, "premerge": premerge}

    def decode_masks(self, embeds, embeds_info, input_ids, image, boxes):
        # Teacher-forced pass over prompt + response: the <OBJ> hidden states (last layer, before the final norm) seed the mask queries.
        position_ids, visual_pos_masks, deepstack = self.build_image_inputs(embeds, embeds_info)
        _, hidden = self.model(None, embeds=embeds, intermediate_output=-1, final_layer_norm_intermediate=False, position_ids=position_ids,
                               embeds_info=embeds_info, visual_pos_masks=visual_pos_masks, deepstack_embeds=deepstack)
        info = next(e for e in embeds_info if e.get("type") == "image")
        # The image placeholder expands to info["size"] embeddings.
        positions = [i + info["size"] - 1 for i, t in enumerate(input_ids) if t == OBJ_TOKEN][:len(boxes)]
        t_seg = hidden[0, positions].float()
        boxes = boxes[:len(positions)]

        gh, gw = (int(v) for v in info["extra"]["grid"][0, 1:])
        m = self.visual.spatial_merge_size
        # Undo Qwen's merge-local [gh/m, gw/m, m, m] patch order.
        features = [f.float().reshape(gh // m, gw // m, m, m, -1).permute(4, 0, 2, 1, 3).reshape(1, -1, gh, gw) for f in info["extra"]["premerge"]]

        cnn_image = F.interpolate(image.movedim(-1, 1)[:, :3].to(embeds.device, torch.float32), size=(CNN_SIZE, CNN_SIZE), mode="bicubic", antialias=True, align_corners=False).clamp(0.0, 1.0)
        boxes = torch.tensor(boxes, device=embeds.device, dtype=torch.float32)
        alphas = [self.mask_decoder(cnn_image, features, t_seg[i:i + 1], boxes[i:i + 1]).sigmoid() for i in range(t_seg.shape[0])]
        return F.interpolate(torch.cat(alphas)[None], size=image.shape[1:3], mode="bilinear", align_corners=False)[0]


class WeVisOneClipModel(Qwen3VLClipModel):
    def __init__(self, device="cpu", layer="hidden", layer_idx=-1, dtype=None, attention_mask=True, model_options={}):
        sd1_clip.SDClipModel.__init__(self, device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"pad": 151643},
                                      layer_norm_hidden_state=False, model_class=WeVisOneQwen3VL, enable_attention_masks=attention_mask,
                                      return_attention_masks=attention_mask, model_options=model_options)

    def decode_masks(self, tokens, generated_ids, image, boxes):
        input_ids = [t[0] for t in next(iter(tokens.values()))[0]] + list(generated_ids)
        embeds, _, _, embeds_info = self.process_tokens([input_ids], self.execution_device)
        dtype = torch.bfloat16 if comfy.model_management.should_use_bf16(self.execution_device) else torch.float32
        return self.transformer.decode_masks(embeds.to(dtype), embeds_info, input_ids, image, boxes)


class WeVisOneTEModel(sd1_clip.SD1ClipModel):
    def __init__(self, device="cpu", dtype=None, model_options={}):
        super().__init__(device=device, dtype=dtype, name="qwen3vl_4b", clip_model=WeVisOneClipModel, model_options=model_options)

    def decode_masks(self, tokens, generated_ids, image, boxes):
        return self.qwen3vl_4b.decode_masks(tokens, generated_ids, image, boxes)


class WeVisOneTokenizer(Qwen3VLTokenizer):
    def __init__(self, embedding_directory=None, tokenizer_data={}):
        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, model_type="qwen3vl_4b")
        self.qwen3vl_4b.tokenizer.add_tokens(ADDED_TOKENS)