wehub-resource-sync
45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from typing import Any
|
|
|
|
from transformers import PretrainedConfig
|
|
from transformers.models.gemma4.configuration_gemma4 import Gemma4VisionConfig
|
|
|
|
|
|
def _init_text_config(self: PretrainedConfig, **kwargs: Any) -> None:
|
|
PretrainedConfig.__init__(self, **kwargs)
|
|
# DiffusionGemma always uses MoE and K=V sharing for full_attention
|
|
# layers. The HF reference removed these config fields entirely.
|
|
if getattr(self, "num_experts", None):
|
|
self.enable_moe_block = True
|
|
self.attention_k_eq_v = True
|
|
|
|
|
|
class DiffusionGemmaTextConfig(PretrainedConfig):
|
|
model_type = "diffusion_gemma_text"
|
|
|
|
def __init__(self, **kwargs: Any):
|
|
_init_text_config(self, **kwargs)
|
|
|
|
|
|
class DiffusionGemmaConfig(PretrainedConfig):
|
|
model_type = "diffusion_gemma"
|
|
|
|
def __init__(
|
|
self,
|
|
text_config: dict[str, Any] | None = None,
|
|
canvas_length: int = 256,
|
|
self_conditioning_size: int | None = None,
|
|
**kwargs: Any,
|
|
):
|
|
self.text_config = DiffusionGemmaTextConfig(**(text_config or {}))
|
|
self.canvas_length = canvas_length
|
|
self.self_conditioning_size = self_conditioning_size
|
|
vision_config = kwargs.pop("vision_config", None)
|
|
if isinstance(vision_config, dict):
|
|
self.vision_config = Gemma4VisionConfig(**vision_config)
|
|
else:
|
|
self.vision_config = vision_config
|
|
self.audio_config = None
|
|
PretrainedConfig.__init__(self, **kwargs)
|