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

class M(torch.nn.Module):
    @property
    def device(self): return next(self.parameters()).device
    @property
    def dtype(self): return next(self.parameters()).dtype

class Tok:
    def __call__(self, prompt, max_length=None, **_):
        prompt = [prompt] if isinstance(prompt, str) else prompt
        ids = torch.zeros(len(prompt), max_length, dtype=torch.long)
        for i, p in enumerate(prompt):
            ids[i, 0] = sum(map(ord, p)) % 100 + 1 if p else 0
        return SimpleNamespace(input_ids=ids)

class Text(M):
    def __init__(self): super().__init__(); self.p = torch.nn.Parameter(torch.ones(()))
    def forward(self, ids): return (ids.float().unsqueeze(-1).repeat(1, 1, 32) / 100,)

class UNet(M):
    def __init__(self):
        super().__init__(); self.p = torch.nn.Parameter(torch.ones(()))
        self.config = SimpleNamespace(in_channels=4, sample_size=8)
    def forward(self, sample, timestep, encoder_hidden_states=None):
        v = encoder_hidden_states[:, 0, 0].view(-1, 1, 1, 1).to(sample)
        return SimpleNamespace(sample=v.expand_as(sample))

class VAE(M):
    def __init__(self):
        super().__init__(); self.p = torch.nn.Parameter(torch.ones(()))
        self.config = SimpleNamespace(block_out_channels=(1,1,1,1), scaling_factor=1.0)
    def decode(self, latents):
        return SimpleNamespace(sample=latents[:, :3].repeat_interleave(8, -1).repeat_interleave(8, -2))

pipe = LDMTextToImagePipeline(VAE(), Text(), Tok(), UNet(), DDIMScheduler())
pipe.set_progress_bar_config(disable=True)
latents = torch.zeros(1, 4, 8, 8)
run = lambda p, gs: pipe(p, latents=latents.clone(), num_inference_steps=1, guidance_scale=gs, output_type="np").images
print(abs(run("cat", 0.0) - run("dog", 0.0)).max())  # 0.0: prompt ignored
print(abs(run("cat", 1.0) - run("dog", 1.0)).max())  # prompt affects output

do_classifier_free_guidance = guidance_scale > 1.0
if do_classifier_free_guidance:
    ...
if do_classifier_free_guidance:
    latents_input = torch.cat([latents] * 2)
    context = torch.cat([negative_prompt_embeds, prompt_embeds])
else:
    latents_input = latents
    context = prompt_embeds

# Same tiny component pattern as above, but VAE has scale factor 2.
# The pipeline requests default height 16, but creates latents at height//8 and decodes to 4.
print(pipe.vae_scale_factor)
print(pipe.unet.config.sample_size * pipe.vae_scale_factor)
print(pipe("x", num_inference_steps=1, guidance_scale=1.0, output_type="np").images.shape)
# observed shape: (1, 4, 4, 3), expected height/width: 16

if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0:
    raise ValueError(
        f"`height` and `width` have to be divisible by {self.vae_scale_factor} but are {height} and {width}."
    )

latents_shape = (
    batch_size,
    self.unet.config.in_channels,
    height // self.vae_scale_factor,
    width // self.vae_scale_factor,
)

import warnings
# Build the same tiny pipeline as Issue 1, but use LMSDiscreteScheduler.
from diffusers import LMSDiscreteScheduler
pipe.scheduler = LMSDiscreteScheduler()
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    pipe("x", num_inference_steps=4, guidance_scale=1.0, output_type="np")
print([str(w.message) for w in caught if "scale_model_input" in str(w.message)])

self.scheduler.set_timesteps(num_inference_steps, device=self._execution_device)
...
latents = latents * self.scheduler.init_noise_sigma
...
latents_input = self.scheduler.scale_model_input(latents_input, t)
noise_pred = self.unet(latents_input, t, encoder_hidden_states=context).sample

from diffusers import DDIMScheduler, LDMSuperResolutionPipeline, UNet2DModel, VQModel

unet = UNet2DModel(sample_size=32, in_channels=6, out_channels=3, block_out_channels=(32, 64))
vqvae = VQModel(in_channels=3, out_channels=3, latent_channels=3, block_out_channels=(32, 64))
pipe = LDMSuperResolutionPipeline(vqvae=vqvae, unet=unet, scheduler=DDIMScheduler())

try:
    pipe.enable_model_cpu_offload(device="cpu")
except Exception as e:
    print(type(e).__name__, e)

class LDMSuperResolutionPipeline(DiffusionPipeline):
    model_cpu_offload_seq = "unet->vqvae"

    ...
    device = self._execution_device
    latents = randn_tensor(latents_shape, generator=generator, device=device, dtype=latents_dtype)
    image = image.to(device=device, dtype=latents_dtype)
    self.scheduler.set_timesteps(num_inference_steps, device=device)

from diffusers.pipelines.latent_diffusion import LDMBertModel
print(LDMBertModel.__name__)

from diffusers.pipelines.latent_diffusion import LDMBertConfig
# ImportError: cannot import name 'LDMBertConfig'

_import_structure["pipeline_latent_diffusion"] = ["LDMBertConfig", "LDMBertModel", "LDMTextToImagePipeline"]

...
from .pipeline_latent_diffusion import LDMBertConfig, LDMBertModel, LDMTextToImagePipeline

from pathlib import Path

for path in [
    "tests/pipelines/latent_diffusion/test_latent_diffusion.py",
    "tests/pipelines/latent_diffusion/test_latent_diffusion_superresolution.py",
]:
    text = Path(path).read_text()
    print(path, "@slow" in text, "@nightly" in text)

from ...testing_utils import slow

@slow
@require_torch_accelerator
class LDMTextToImagePipelineSlowTests(unittest.TestCase):
    ...

@slow
@require_torch
class LDMSuperResolutionPipelineSlowTests(unittest.TestCase):
    ...