import torch
from types import SimpleNamespace
from diffusers import DDIMScheduler, KandinskyPipeline

class Module(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.dummy = torch.nn.Parameter(torch.zeros(()))
    @property
    def dtype(self): return self.dummy.dtype
    @property
    def device(self): return self.dummy.device

class Tokenizer:
    model_max_length = 77
    def __call__(self, prompt, **kwargs):
        b = len(prompt) if isinstance(prompt, list) else 1
        return SimpleNamespace(input_ids=torch.ones(b, 4, dtype=torch.long), attention_mask=torch.ones(b, 4, dtype=torch.long))
    def batch_decode(self, ids): return [""]

class TextEncoder(Module):
    def forward(self, input_ids, attention_mask):
        b, s = input_ids.shape
        return torch.zeros(b, 32), torch.zeros(b, s, 32)

class UNet(Module):
    config = SimpleNamespace(in_channels=4)
    def forward(self, sample, timestep, encoder_hidden_states, added_cond_kwargs, return_dict=False):
        assert added_cond_kwargs["image_embeds"].shape[0] == sample.shape[0], (added_cond_kwargs["image_embeds"].shape, sample.shape)
        return (torch.zeros(sample.shape[0], 8, sample.shape[2], sample.shape[3]),)

class Movq(Module):
    config = SimpleNamespace(block_out_channels=[1])
    def decode(self, latents, force_not_quantize=True): return {"sample": latents[:, :3]}

pipe = KandinskyPipeline(TextEncoder(), Tokenizer(), UNet(), DDIMScheduler(num_train_timesteps=2), Movq()).to("cpu")
pipe.set_progress_bar_config(disable=True)
pipe(prompt="p", image_embeds=torch.zeros(1, 32), negative_image_embeds=torch.zeros(1, 32), guidance_scale=1.0, num_images_per_prompt=2, num_inference_steps=1, output_type="pt")

image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0).to(
    dtype=prompt_embeds.dtype, device=device
)

if do_classifier_free_guidance:
    negative_image_embeds = negative_image_embeds.repeat_interleave(num_images_per_prompt, dim=0).to(
        dtype=prompt_embeds.dtype, device=device
    )
    image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0)

import numpy as np
import torch
from PIL import Image
from types import SimpleNamespace
from diffusers import KandinskyImg2ImgCombinedPipeline

class Prior:
    def __call__(self, **kwargs):
        return (torch.zeros(2, 32), torch.zeros(2, 32))

class Decoder:
    def __call__(self, prompt, image, **kwargs):
        assert isinstance(image, list) and len(image) == 2, (type(image), image if isinstance(image, list) else None)
        return SimpleNamespace(images="ok")

pipe = object.__new__(KandinskyImg2ImgCombinedPipeline)
pipe.prior_pipe = Prior()
pipe.decoder_pipe = Decoder()
pipe.maybe_free_model_hooks = lambda: None

image = Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8))
pipe(prompt="horse", image=image, num_images_per_prompt=2)

prompt = [prompt] if not isinstance(prompt, (list, tuple)) else prompt
image = [image] if isinstance(image, PIL.Image.Image) else image

import torch
from diffusers import Kandinsky5Transformer3DModel

model = Kandinsky5Transformer3DModel(
    in_visual_dim=4, in_text_dim=8, in_text_dim2=4, time_dim=8, out_visual_dim=4,
    patch_size=(1, 1, 1), model_dim=8, ff_dim=16, num_text_blocks=1,
    num_visual_blocks=1, axes_dims=(2, 2, 4), visual_cond=False,
).eval().to(dtype=torch.bfloat16)

model(
    hidden_states=torch.randn(1, 1, 2, 2, 4, dtype=torch.bfloat16),
    encoder_hidden_states=torch.randn(1, 3, 8, dtype=torch.bfloat16),
    timestep=torch.tensor([1], dtype=torch.bfloat16),
    pooled_projections=torch.randn(1, 4, dtype=torch.bfloat16),
    visual_rope_pos=[torch.arange(1), torch.arange(2), torch.arange(2)],
    text_rope_pos=torch.arange(3),
    return_dict=False,
)

time_embedder_dtype = self.in_layer.weight.dtype
if time_embed.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8:
    time_embed = time_embed.to(time_embedder_dtype)
time_embed = self.out_layer(self.activation(self.in_layer(time_embed)))

orig_dtype = x.dtype
x_ = x.reshape(*x.shape[:-1], -1, 1, 2).float()
x_out = (rope * x_).sum(dim=-1)
return x_out.reshape(*x.shape).to(orig_dtype)

from diffusers import Kandinsky5Transformer3DModel

print(getattr(Kandinsky5Transformer3DModel, "_no_split_modules", None))
print(Kandinsky5Transformer3DModel._repeated_blocks)
assert Kandinsky5Transformer3DModel._no_split_modules is not None

_no_split_modules = ["Kandinsky5TransformerEncoderBlock", "Kandinsky5TransformerDecoderBlock"]
_repeated_blocks = ["Kandinsky5TransformerEncoderBlock", "Kandinsky5TransformerDecoderBlock"]

from pathlib import Path

for path in sorted(Path("tests/pipelines/kandinsky").glob("test_*.py")):
    text = path.read_text()
    print(path.as_posix(), "slow=", "@slow" in text, "nightly=", "@nightly" in text, "xfail=", "xfail" in text)

has_model_test = any(
    "Kandinsky5Transformer3DModel" in path.read_text(errors="ignore")
    for path in Path("tests/models").rglob("test_*.py")
)
print("Kandinsky5 model test:", has_model_test)

# Fast
# - no-CFG num_images_per_prompt > 1 for KandinskyPipeline and KandinskyInpaintPipeline
# - single PIL image + num_images_per_prompt > 1 for combined img2img/inpaint
# - Kandinsky5Transformer3DModel .to(torch.bfloat16) forward

# Slow
# - KandinskyPriorPipeline
# - KandinskyCombinedPipeline
# - KandinskyImg2ImgCombinedPipeline
# - KandinskyInpaintCombinedPipeline
# - @slow inpaint coverage, or mark the existing nightly integration as slow too