import json
import torch
from diffusers import AutoGuidance, SkipLayerGuidance, ClassifierFreeGuidance, SmoothedEnergyGuidance, SmoothedEnergyGuidanceConfig, AdaptiveProjectedGuidance, PerturbedAttentionGuidance, ClassifierFreeZeroStarGuidance, TangentialClassifierFreeGuidance, LayerSkipConfig
from diffusers.modular_pipelines import SequentialPipelineBlocks, ComponentSpec
from diffusers.modular_pipelines.stable_diffusion_xl import TEXT2IMAGE_BLOCKS
from diffusers.utils.logging import set_verbosity_debug

set_verbosity_debug()

blocks = SequentialPipelineBlocks.from_blocks_dict(TEXT2IMAGE_BLOCKS)

modular_repo_id = "YiYiXu/modular-loader-t2i"
pipeline = blocks.init_pipeline(modular_repo_id)
pipeline.load_default_components(torch_dtype=torch.float16)
pipeline.to("cuda")

for guider_cls in [
    AutoGuidance,
    SkipLayerGuidance,
    ClassifierFreeGuidance,
    SmoothedEnergyGuidance,
    AdaptiveProjectedGuidance,
    PerturbedAttentionGuidance,
    ClassifierFreeZeroStarGuidance,
    TangentialClassifierFreeGuidance
]:
    print(f"Testing {guider_cls.__name__}...")
    
    kwargs = {"guidance_scale": 5.0}
    if guider_cls is AutoGuidance:
        kwargs.update({"auto_guidance_config": LayerSkipConfig(indices=[9, 10], fqn="mid_block.attentions.0.transformer_blocks", skip_attention=True, skip_ff=True, dropout=0.75)})
        kwargs.update({"stop": 0.6})
    elif guider_cls is SkipLayerGuidance:
        kwargs.update({"skip_layer_config": LayerSkipConfig(indices=[9, 10], fqn="mid_block.attentions.0.transformer_blocks", skip_attention=True, skip_ff=True)})
        kwargs.update({"skip_layer_guidance_stop": 0.4})
    elif guider_cls is SmoothedEnergyGuidance:
        kwargs.update({"seg_guidance_config": SmoothedEnergyGuidanceConfig(indices=[9, 10], fqn="mid_block.attentions.0.transformer_blocks")})
        kwargs.update({"seg_guidance_stop": 0.8})
    elif guider_cls is PerturbedAttentionGuidance:
        kwargs.update({"perturbed_guidance_config": LayerSkipConfig(indices=[9, 10], fqn="mid_block.attentions.0.transformer_blocks", skip_attention=False, skip_attention_scores=True, skip_ff=False)})
        kwargs.update({"perturbed_guidance_stop": 0.5})

    pipeline.update_components(
        guider=ComponentSpec(
            name="cfg",
            type_hint=guider_cls,
            config=kwargs,
            default_creation_method="from_config",
        )
    )
    pipeline.save_pretrained("test_modular_t2i_guider")
    pipeline.components["guider"].save_config("test_modular_t2i_guider/guider")
    
    with open("test_modular_t2i_guider/guider/guider_config.json", "r") as f:
        guider_config = json.load(f)
    guider_config["guidance_scale"] = 10.0
    with open("test_modular_t2i_guider/guider/guider_config.json", "w") as f:
        json.dump(guider_config, f, indent=2)
    
    guider_spec = ComponentSpec(
        name="cfg",
        type_hint=guider_cls,
        config=guider_cls.from_config("test_modular_t2i_guider/guider").config,
        default_creation_method="from_config",
    )
    pipeline.update_components(guider=guider_spec.create())
    print(pipeline.components["guider"])

    image = pipeline(prompt="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", num_inference_steps=20, output="images", generator=torch.Generator().manual_seed(42))[0]
    output_filename = f"output_guider_{guider_cls.__name__.lower()}.png"
    image.save(output_filename)