# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""CPU-only unit tests for the SDXL diffusion family.

SDXL is the one U-Net family: the denoiser is ``pipe.unet`` (not ``pipe.transformer``)
and a single-file ``.safetensors`` is the whole pipeline (not a transformer-only file).
These tests cover the pure helpers that encode those differences -- family detection,
the ``denoiser_attr`` / ``single_file_is_pipeline`` flags, the non-GGUF trust allowlist,
the VAE-dtype alignment reading the U-Net denoiser, and the LoRA-support gate -- with no
torch/diffusers/GPU needed.
"""

from __future__ import annotations

import types

import pytest

from core.inference import diffusion_lora
from core.inference.diffusion import (
    DiffusionBackend,
    _is_trusted_diffusion_repo,
    resolve_model_kind,
)
from core.inference.diffusion_families import detect_family, family_sd_cpp_supported


def test_sdxl_family_shape():
    fam = detect_family("stabilityai/stable-diffusion-xl-base-1.0")
    assert fam is not None and fam.name == "sdxl"
    assert fam.pipeline_class == "StableDiffusionXLPipeline"
    # The denoiser is a U-Net, addressed via pipe.unet (DiT families use pipe.transformer).
    assert fam.denoiser_attr == "unet"
    assert fam.transformer_class == "UNet2DConditionModel"
    # A single-file SDXL checkpoint is the whole pipeline, loaded via the pipeline class.
    assert fam.single_file_is_pipeline is True
    # Image-conditioned + ControlNet workflows are the standard SDXL pipelines.
    assert fam.img2img_pipeline_class == "StableDiffusionXLImg2ImgPipeline"
    assert fam.inpaint_pipeline_class == "StableDiffusionXLInpaintPipeline"
    assert fam.controlnet_pipeline_class == "StableDiffusionXLControlNetPipeline"
    assert fam.controlnet_model_class == "ControlNetModel"
    # Real CFG; SDXL uses guidance_scale, not a distilled true_cfg_scale.
    assert fam.cfg_kwarg == "guidance_scale"


def test_sdxl_detection_by_repo_and_override():
    assert detect_family("stabilityai/sdxl-turbo").name == "sdxl"
    assert detect_family("some-org/My-Cool-SDXL-Merge").name == "sdxl"
    assert detect_family("some-org/stable-diffusion-xl-anime").name == "sdxl"
    assert detect_family("x", override = "sdxl").name == "sdxl"
    # A GGUF DiT family must NOT be swallowed by the SDXL match.
    assert detect_family("unsloth/FLUX.1-schnell-GGUF").name == "flux.1"


def test_dit_families_keep_transformer_denoiser():
    # The generalisation must not change existing DiT families: they stay on
    # pipe.transformer and their single file is transformer-only.
    for rid in ("unsloth/FLUX.1-schnell-GGUF", "unsloth/Qwen-Image-GGUF", "unsloth/Z-Image-GGUF"):
        fam = detect_family(rid)
        assert fam.denoiser_attr == "transformer"
        assert fam.single_file_is_pipeline is False


def test_sdxl_has_no_native_sd_cpp_mapping():
    # No single-file VAE/TE mapping yet, so the no-GPU route falls back to diffusers
    # rather than trying to drive sd-cli.
    assert family_sd_cpp_supported(detect_family("stabilityai/sdxl-turbo")) is False


def test_sdxl_base_repos_are_trusted_non_gguf():
    # Official safetensors-only base repos are allowlisted so their catalog entries load.
    assert _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-base-1.0")
    assert _is_trusted_diffusion_repo("stabilityai/sdxl-turbo")
    # The refiner is img2img-only and is intentionally NOT allowlisted (see
    # test_sdxl_refiner_not_trusted).
    # Case-insensitive match.
    assert _is_trusted_diffusion_repo("StabilityAI/SDXL-Turbo")
    # A random repo (even one that detects as SDXL) is NOT trusted for a non-GGUF load.
    assert not _is_trusted_diffusion_repo("randomorg/my-sdxl-merge")
    assert not _is_trusted_diffusion_repo("stabilityai/sdxl-turbo-evil")


def test_sdxl_model_kind_resolution():
    # A full-pipeline load (no single-file name) is "pipeline"; a single .safetensors
    # is "single_file" (handled by the whole-pipeline branch for SDXL).
    assert resolve_model_kind(None) == "pipeline"
    assert resolve_model_kind("sdxl.safetensors") == "single_file"


class _FakeVae:
    def __init__(self, dtype):
        self._dtype = dtype
        self.moved_to = None

    def parameters(self):
        yield types.SimpleNamespace(dtype = self._dtype)

    def to(self, dtype = None):
        self.moved_to = dtype
        self._dtype = dtype


def test_align_vae_dtype_uses_unet_denoiser():
    # For SDXL the denoiser lives at pipe.unet; _align_vae_dtype must read it (a pipe
    # with only .unet and no .transformer) and cast the VAE to the U-Net's dtype. The
    # dtype is read from a parameter (denoiser has no .dtype), so use a _FakeVae denoiser.
    import torch

    vae = _FakeVae(dtype = torch.float32)
    unet = _FakeVae(dtype = torch.bfloat16)
    pipe = types.SimpleNamespace(unet = unet, vae = vae)
    DiffusionBackend._align_vae_dtype(pipe, "unet")
    assert vae.moved_to == torch.bfloat16


def test_align_vae_dtype_transformer_default_unchanged():
    # DiT default: reads pipe.transformer; a pipe with no transformer is a safe no-op.
    import torch

    vae = _FakeVae(dtype = torch.float32)
    transformer = _FakeVae(dtype = torch.bfloat16)
    pipe = types.SimpleNamespace(transformer = transformer, vae = vae)
    DiffusionBackend._align_vae_dtype(pipe)
    assert vae.moved_to == torch.bfloat16
    # No denoiser attribute -> no-op (does not raise, does not move the VAE).
    vae2 = _FakeVae(dtype = torch.float32)
    DiffusionBackend._align_vae_dtype(types.SimpleNamespace(vae = vae2), "unet")
    assert vae2.moved_to is None


def test_align_vae_dtype_skips_gguf_packed_uint8_params():
    # A GGUF-quantized transformer's leading parameters are packed uint8 storage; the
    # dtype probe must skip them and use the first FLOATING dtype, or nn.Module.to()
    # rejects the integer dtype and an Edit/img2img call 500s (regression: Qwen-Image-
    # Edit GGUF). All-integer params (no floating dtype at all) must be a clean no-op.
    import torch

    class _GgufDenoiser:
        def parameters(self):
            yield types.SimpleNamespace(dtype = torch.uint8)  # packed GGUF block
            yield types.SimpleNamespace(dtype = torch.bfloat16)  # compute dtype

    vae = _FakeVae(dtype = torch.float32)
    pipe = types.SimpleNamespace(transformer = _GgufDenoiser(), vae = vae)
    DiffusionBackend._align_vae_dtype(pipe)
    assert vae.moved_to == torch.bfloat16

    class _AllPacked:
        def parameters(self):
            yield types.SimpleNamespace(dtype = torch.uint8)

    vae2 = _FakeVae(dtype = torch.float32)
    DiffusionBackend._align_vae_dtype(types.SimpleNamespace(transformer = _AllPacked(), vae = vae2))
    assert vae2.moved_to is None


def test_sdxl_lora_supported_on_diffusers():
    # SDXL is bf16/bnb-4bit on diffusers -> LoRA is allowed (unlike GGUF-via-diffusers).
    assert diffusion_lora.supports_lora(
        engine = "diffusers", family = "sdxl", model_kind = "pipeline", transformer_quant = None
    )
    assert diffusion_lora.supports_lora(
        engine = "diffusers", family = "sdxl", model_kind = "single_file", transformer_quant = None
    )


def test_pipeline_prefetch_skips_non_torch_artifacts():
    # The SDXL Base repo ships fp16 variants, ONNX, OpenVINO and Flax exports next to
    # the default safetensors; from_pretrained (no variant kwarg) loads only the
    # default torch weights, so the prefetch filter must skip everything else or a
    # catalog load pulls tens of GB of unused artifacts.
    from core.inference.diffusion import _pipeline_file_downloaded as keep

    assert keep("model_index.json")
    assert keep("unet/diffusion_pytorch_model.safetensors")
    assert keep("text_encoder/model.safetensors")
    assert keep("scheduler/scheduler_config.json")
    assert not keep("sd_xl_base_1.0.safetensors")  # top-level single-file twin
    assert not keep("unet/diffusion_pytorch_model.fp16.safetensors")
    assert not keep("text_encoder/model.onnx")
    assert not keep("text_encoder/openvino_model.bin")
    assert not keep("unet/flax_model.msgpack")
    assert not keep("vae_decoder/model.onnx_data")
    assert not keep("assets/preview.png")


def test_sdxl_refiner_not_trusted():
    # The refiner is an img2img-only pipeline; the sdxl family loads every repo as the
    # base txt2img pipeline, so the refiner must NOT be allowlisted for a non-GGUF load.
    assert not _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-refiner-1.0")
    # The base and turbo remain trusted.
    assert _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-base-1.0")
    assert _is_trusted_diffusion_repo("stabilityai/sdxl-turbo")


def test_sdxl_gguf_load_rejected_up_front():
    # SDXL has no transformer-only GGUF variant (its single file is the whole pipeline),
    # so a GGUF request must fail cheap validation before the GPU handoff.
    backend = DiffusionBackend()
    with pytest.raises(ValueError, match = "no GGUF"):
        backend.validate_load_request(
            "some-org/my-sdxl.gguf", gguf_filename = "my-sdxl.gguf", family_override = "sdxl"
        )


def test_base_config_filter_skips_weights():
    # For a whole-pipeline single file, the base repo supplies only config/tokenizer, not
    # its (unused) weight tensors.
    from core.inference.diffusion import _base_config_file_downloaded as keep

    assert keep("model_index.json")
    assert keep("text_encoder/config.json")
    assert keep("tokenizer/vocab.json")
    assert keep("scheduler/scheduler_config.json")
    assert not keep("unet/diffusion_pytorch_model.safetensors")
    assert not keep("vae/diffusion_pytorch_model.bin")
    assert not keep("text_encoder/model.onnx")
    # transformer/ and assets/ stay excluded (inherited from _base_file_downloaded).
    assert not keep("transformer/config.json")
    assert not keep("assets/x.png")