from diffusers.utils import is_flax_available, is_transformers_available

print(is_flax_available(), is_transformers_available())

from diffusers import FlaxStableDiffusionXLPipeline
print("top-level:", FlaxStableDiffusionXLPipeline)

from diffusers.pipelines.stable_diffusion_xl import FlaxStableDiffusionXLPipeline
print("subpackage:", FlaxStableDiffusionXLPipeline)

try:
    if not (is_transformers_available() and is_flax_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from ...utils import dummy_flax_and_transformers_objects

    _dummy_objects.update(get_objects_from_module(dummy_flax_and_transformers_objects))
else:
    from ...schedulers.scheduling_pndm_flax import PNDMSchedulerState

    _additional_imports.update({"PNDMSchedulerState": PNDMSchedulerState})
    _import_structure["pipeline_flax_stable_diffusion_xl"] = ["FlaxStableDiffusionXLPipeline"]

from types import SimpleNamespace
import torch
from diffusers import StableDiffusionXLImg2ImgPipeline

class FakePipe:
    pass

pipe = FakePipe()
pipe.config = SimpleNamespace(requires_aesthetics_score=False)
pipe.unet = SimpleNamespace(
    config=SimpleNamespace(addition_time_embed_dim=1),
    add_embedding=SimpleNamespace(linear_1=SimpleNamespace(in_features=7)),
)

_, negative = StableDiffusionXLImg2ImgPipeline._get_add_time_ids(
    pipe,
    original_size=(64, 64),
    crops_coords_top_left=(1, 2),
    target_size=(64, 64),
    aesthetic_score=6.0,
    negative_aesthetic_score=2.0,
    negative_original_size=(32, 32),
    negative_crops_coords_top_left=(9, 10),
    negative_target_size=(32, 32),
    dtype=torch.float32,
    text_encoder_projection_dim=1,
)
print(negative.tolist())  # contains 1, 2; expected 9, 10

add_neg_time_ids = list(negative_original_size + negative_crops_coords_top_left + negative_target_size)

import torch
from PIL import Image
from diffusers import StableDiffusionXLInpaintPipeline

pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
    "hf-internal-testing/tiny-stable-diffusion-xl-inpaint-pipe",
    add_watermarker=False,
)
pipe.set_progress_bar_config(disable=True)

called = {"cleanup": False}
pipe.maybe_free_model_hooks = lambda: called.__setitem__("cleanup", True)

out = pipe(
    "a cat",
    image=Image.new("RGB", (64, 64), "white"),
    mask_image=Image.new("L", (64, 64), 0),
    num_inference_steps=1,
    strength=1.0,
    output_type="latent",
    return_dict=False,
    generator=torch.Generator("cpu").manual_seed(0),
)

print(type(out).__name__, called)  # StableDiffusionXLPipelineOutput {'cleanup': False}

else:
    image = latents

self.maybe_free_model_hooks()

if not return_dict:
    return (image,)

return StableDiffusionXLPipelineOutput(images=image)

import torch
from PIL import Image
from diffusers import StableDiffusionXLImg2ImgPipeline

class SentinelWatermark:
    def apply_watermark(self, images):
        raise RuntimeError(f"watermark called for {tuple(images.shape)}")

pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    "hf-internal-testing/tiny-stable-diffusion-xl-pipe",
    add_watermarker=False,
)
pipe.set_progress_bar_config(disable=True)
pipe.watermark = SentinelWatermark()

pipe(
    "a cat",
    image=Image.new("RGB", (64, 64), "white"),
    strength=1.0,
    num_inference_steps=1,
    output_type="latent",
    generator=torch.Generator("cpu").manual_seed(0),
)

if not output_type == "latent":
    if self.watermark is not None:
        image = self.watermark.apply_watermark(image)
    image = self.image_processor.postprocess(image, output_type=output_type)

from types import SimpleNamespace
import torch
from diffusers.modular_pipelines.stable_diffusion_xl.encoders import StableDiffusionXLInpaintVaeEncoderStep

class FakeVAE:
    config = SimpleNamespace(
        force_upcast=False,
        latents_mean=[0.0, 0.0, 0.0, 0.0],
        latents_std=[1.0, 1.0, 1.0, 1.0],
        scaling_factor=0.18215,
    )

    def encode(self, image):
        return SimpleNamespace(latents=torch.ones(image.shape[0], 4, 2, 2, dtype=image.dtype))

components = SimpleNamespace(vae=FakeVAE())
StableDiffusionXLInpaintVaeEncoderStep()._encode_vae_image(
    components, torch.zeros(1, 3, 16, 16), generator=None
)

image_latents = (image_latents - latents_mean) * components.vae.config.scaling_factor / latents_std

from diffusers import StableDiffusionXLAutoBlocks

doc = StableDiffusionXLAutoBlocks.__doc__ or ""
print(doc.count("TODO: Add description."))
assert "TODO: Add description." not in doc

from pathlib import Path

paths = [
    Path("tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py"),
    Path("tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_img2img.py"),
    Path("tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_inpaint.py"),
    Path("tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_instruction_pix2pix.py"),
    Path("tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py"),
]

for path in paths:
    print(path, path.read_text(encoding="utf-8").count("@slow"))

print("Flax SDXL tests:", list(Path("tests").rglob("*flax*sdxl*")) + list(Path("tests").rglob("*sdxl*flax*")))