import torch
from diffusers import ChromaTransformer2DModel

torch.manual_seed(0)
model = ChromaTransformer2DModel(
    in_channels=4, out_channels=4, num_layers=1, num_single_layers=1,
    attention_head_dim=4, num_attention_heads=1, joint_attention_dim=8,
    axes_dims_rope=(0, 2, 2), approximator_num_channels=16,
    approximator_hidden_dim=8, approximator_layers=1,
).eval()

base_encoder = torch.randn(1, 3, 8)
changed_encoder = base_encoder.clone()
changed_encoder[:, 2] += 1000  # token 2 is masked

common = dict(
    hidden_states=torch.randn(1, 2, 4),
    timestep=torch.ones(1),
    txt_ids=torch.zeros(3, 3),
    img_ids=torch.zeros(2, 3),
)

for dtype in (torch.bool, torch.float32):
    mask = torch.tensor([[1, 1, 0, 1, 1]], dtype=dtype)
    with torch.no_grad():
        a = model(encoder_hidden_states=base_encoder, attention_mask=mask, **common).sample
        b = model(encoder_hidden_states=changed_encoder, attention_mask=mask, **common).sample
    print(dtype, (a - b).abs().max().item())

# In all Chroma _get_t5_prompt_embeds methods:
attention_mask = mask_indices <= seq_lengths.unsqueeze(1)

# In all Chroma _prepare_attention_mask methods:
attention_mask = attention_mask.to(dtype=torch.bool)
image_attention_mask = torch.ones(
    batch_size, sequence_length, device=attention_mask.device, dtype=torch.bool
)
attention_mask = torch.cat([attention_mask, image_attention_mask], dim=1)

from diffusers import ChromaTransformer2DModel

transformer = ChromaTransformer2DModel(
    in_channels=4, out_channels=4, num_layers=0, num_single_layers=0,
    attention_head_dim=4, num_attention_heads=1, joint_attention_dim=8,
    axes_dims_rope=(0, 2, 2), approximator_num_channels=16,
    approximator_hidden_dim=8,
)

print("guidance_embeds" in transformer.config)
print(transformer.config.guidance_embeds)  # AttributeError

# Remove the inpaint-only guidance_embeds block entirely.
# ChromaTransformer2DModel has no guidance input and the local `guidance` value is unused.

import torch
from diffusers import ChromaTransformer2DModel

torch.manual_seed(0)
model = ChromaTransformer2DModel(
    in_channels=4, out_channels=4, num_layers=0, num_single_layers=1,
    attention_head_dim=4, num_attention_heads=1, joint_attention_dim=8,
    axes_dims_rope=(0, 2, 2), approximator_num_channels=16,
    approximator_hidden_dim=8, approximator_layers=1,
).eval()

inputs = dict(
    hidden_states=torch.randn(1, 2, 4, requires_grad=True),
    encoder_hidden_states=torch.randn(1, 2, 8, requires_grad=True),
    timestep=torch.ones(1),
    txt_ids=torch.zeros(2, 3),
    img_ids=torch.zeros(2, 3),
    attention_mask=torch.tensor([[1, 0, 1, 1]], dtype=torch.bool),
)

out_no_ckpt = model(**inputs).sample.detach()
model.enable_gradient_checkpointing()
out_ckpt = model(**inputs).sample.detach()
print((out_no_ckpt - out_ckpt).abs().max().item())

encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
    block, hidden_states, encoder_hidden_states, temb, image_rotary_emb, attention_mask, joint_attention_kwargs
)

hidden_states = self._gradient_checkpointing_func(
    block, hidden_states, temb, image_rotary_emb, attention_mask, joint_attention_kwargs
)

from pathlib import Path

files = sorted(str(p).replace("\\", "/") for p in Path("tests").rglob("*chroma*.py"))
print("\n".join(files))
print("has_inpaint_fast_test", any("inpaint" in f for f in files))
print("has_slow_marker", any("@slow" in Path(f).read_text(encoding="utf-8") for f in files))

# Add tests/pipelines/chroma/test_pipeline_chroma_inpainting.py using the existing
# Chroma tiny components plus image/mask tensors, and add at least one @slow Chroma
# pipeline test against a published Chroma checkpoint or saved test slices.

import warnings
from diffusers import ChromaTransformer2DModel

with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    ChromaTransformer2DModel(
        in_channels=4, out_channels=4, num_layers=0, num_single_layers=0,
        attention_head_dim=4, num_attention_heads=1, joint_attention_dim=8,
        axes_dims_rope=(0, 2, 2), approximator_num_channels=16,
        approximator_hidden_dim=8,
    )

print([str(w.message) for w in caught if "FluxPosEmbed" in str(w.message)])

from ..embeddings import PixArtAlphaTextProjection, Timesteps, get_timestep_embedding
from .transformer_flux import FluxAttention, FluxAttnProcessor, FluxPosEmbed

from pathlib import Path

inpaint = Path("src/diffusers/pipelines/chroma/pipeline_chroma_inpainting.py").read_text()
output = Path("src/diffusers/pipelines/chroma/pipeline_output.py").read_text()

for token in ["contentReference[oaicite", "The Flux pipeline", "[`DDIMScheduler`]", "[`CLIPTextModel`]", "text_encoder_2"]:
    print(token, token in inpaint)
print("Stable Diffusion output doc", "Output class for Stable Diffusion pipelines." in output)

# Replace the copied inpaint docstring with Chroma-specific text:
# - Chroma image inpainting
# - FlowMatchEulerDiscreteScheduler
# - T5EncoderModel / T5TokenizerFast
# - no text_encoder_2/tokenizer_2/CLIP text encoder
# Also change ChromaPipelineOutput to "Output class for Chroma pipelines."