import torch
from diffusers.models.attention_processor import Attention
from diffusers.models.transformers.transformer_glm_image import GlmImageAttnProcessor

attn = Attention(query_dim=4, heads=1, dim_head=4, out_dim=4, bias=False, processor=GlmImageAttnProcessor())
with torch.no_grad():
    attn.to_q.weight.zero_()
    attn.to_k.weight.zero_()
    attn.to_v.weight.copy_(torch.eye(4))
    attn.to_out[0].weight.copy_(torch.eye(4))

encoder_hidden_states = torch.tensor([[[0., 0., 0., 0.], [1000., 0., 0., 0.]]])
hidden_states = torch.zeros(1, 1, 4)
attention_mask = torch.tensor([[True, False]])

image_out, _ = attn(
    hidden_states=hidden_states,
    encoder_hidden_states=encoder_hidden_states,
    attention_mask=attention_mask,
)
print(image_out[0, 0, 0].item())  # current: ~155.66; expected: 0 if token 2 is masked

# In GlmImageAttnProcessor.__call__, keep a bool key mask instead of a float QK matrix.
if attention_mask is not None:
    attention_mask = attention_mask.to(device=query.device, dtype=torch.bool)
    cached_seq_length = key.shape[1] - text_seq_length - image_seq_length
    cache_mask = torch.ones((batch_size, cached_seq_length), device=query.device, dtype=torch.bool)
    image_mask = torch.ones((batch_size, image_seq_length), device=query.device, dtype=torch.bool)
    attention_mask = torch.cat([cache_mask, attention_mask, image_mask], dim=1)

import torch
from diffusers import GlmImagePipeline

class Config:
    patch_size = 2

class Transformer:
    config = Config()

pipe = object.__new__(GlmImagePipeline)
pipe.vae_scale_factor = 8
pipe.transformer = Transformer()
pipe._callback_tensor_inputs = ["latents", "prompt_embeds"]

pipe.check_inputs(prompt="x", height=32, width=20, callback_on_step_end_tensor_inputs=["latents"])
latents = pipe.prepare_latents(1, 4, 32, 20, torch.float32, torch.device("cpu"), torch.Generator().manual_seed(0))
print(latents.shape[-1] * pipe.vae_scale_factor)  # 16, not requested width 20

multiple_of = self.vae_scale_factor * self.transformer.config.patch_size * 2
if (height is not None and height % multiple_of != 0) or (width is not None and width % multiple_of != 0):
    raise ValueError(f"`height` and `width` have to be divisible by {multiple_of} but are {height} and {width}.")

import torch
from diffusers import GlmImagePipeline

pipe = object.__new__(GlmImagePipeline)
prompt_embeds = torch.randn(1, 2, 4, dtype=torch.float64)

out, _ = pipe.encode_prompt(
    prompt=None,
    do_classifier_free_guidance=False,
    prompt_embeds=prompt_embeds,
    device=torch.device("cpu"),
    dtype=torch.float32,
)
print(out.dtype, out is prompt_embeds)  # torch.float64 True

if prompt_embeds is None:
    prompt_embeds = self._get_glyph_embeds(prompt, max_sequence_length, device, dtype)
else:
    prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)

if negative_prompt_embeds is not None:
    negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)

if prior_token_ids is not None:
    prior_token_ids = prior_token_ids.to(device=device)
if prior_token_image_ids is not None:
    prior_token_image_ids = [x.to(device=device) for x in prior_token_image_ids]
if source_image_grid_thw is not None:
    source_image_grid_thw = [x.to(device=device) for x in source_image_grid_thw]

import transformers
from diffusers.utils import is_transformers_version
from diffusers.pipelines.glm_image.pipeline_glm_image import GlmImageProcessor, GlmImageForConditionalGeneration
from transformers import ProcessorMixin, PreTrainedModel

print(transformers.__version__)
print(is_transformers_version(">", "4.57.4"))          # True in this env
print(is_transformers_version(">=", "5.0.0.dev0"))     # False
print(hasattr(transformers, "GlmImageProcessor"))      # False
print(GlmImageProcessor is ProcessorMixin, GlmImageForConditionalGeneration is PreTrainedModel)

GLM_IMAGE_TRANSFORMERS_MIN_VERSION = "5.0.0.dev0"  # or the first released transformers version with these classes

if is_transformers_available() and is_transformers_version(">=", GLM_IMAGE_TRANSFORMERS_MIN_VERSION):
    from transformers import GlmImageForConditionalGeneration, GlmImageProcessor
else:
    raise OptionalDependencyNotAvailable()

from pathlib import Path

test_text = Path("tests/pipelines/glm_image/test_glm_image.py").read_text(encoding="utf-8")
model_doc = Path("docs/source/en/api/models/glm_image_transformer2d.md").read_text(encoding="utf-8")

print("@slow present:", "@slow" in test_text)
print("fast class requires accelerator:", "@require_torch_accelerator" in test_text)
print("fast test downloads full GLM processor:", "zai-org/GLM-Image" in test_text)
print("model docs still contain TODO:", "TODO" in model_doc)