Repository: huggingface/diffusers
File: tests/lora/utils.py (Python)

Code before:
    has_three_text_encoders = False
    text_encoder_cls, text_encoder_id, text_encoder_subfolder = None, None, ""
    text_encoder_2_cls, text_encoder_2_id, text_encoder_2_subfolder = None, None, ""
    text_encoder_3_cls, text_encoder_3_id, text_encoder_3_subfolder = None, None, ""
    tokenizer_cls, tokenizer_id, tokenizer_subfolder = None, None, ""
    tokenizer_2_cls, tokenizer_2_id, tokenizer_2_subfolder = None, None, ""
    tokenizer_3_cls, tokenizer_3_id, tokenizer_3_subfolder = None, None, ""

    unet_kwargs = None
    transformer_cls = None
    transformer_kwargs = None
    vae_cls = AutoencoderKL
    vae_kwargs = None

    text_encoder_target_modules = ["q_proj", "k_proj", "v_proj", "out_proj"]
    denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0"]

    cached_non_lora_outputs = {}

    @pytest.fixture(scope="class", autouse=True)
    def cache_non_lora_outputs(self):
        """
        This fixture will be executed once per test class and will populate
        the cached_non_lora_outputs dictionary.
        """
        for scheduler_cls in self.scheduler_classes:
            # Check if the output for this scheduler is already cached to avoid re-running
            if scheduler_cls.__name__ in self.cached_non_lora_outputs:
                continue

            components, _, _ = self.get_dummy_components(scheduler_cls)
            pipe = self.pipeline_class(**components)
            pipe = pipe.to(torch_device)
            pipe.set_progress_bar_config(disable=None)

            # Always ensure the inputs are without the `generator`. Make sure to pass the `generator`
            # explicitly.
            _, _, inputs = self.get_dummy_inputs(with_generator=False)
            output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
            self.cached_non_lora_outputs[scheduler_cls.__name__] = output_no_lora

        # Ensures that there's no inconsistency when reusing the cache.
        yield
        self.cached_non_lora_outputs.clear()

    def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=None):
        if self.unet_kwargs and self.transformer_kwargs:
            raise ValueError("Both `unet_kwargs` and `transformer_kwargs` cannot be specified.")
        if self.has_two_text_encoders and self.has_three_text_encoders:
            raise ValueError("Both `has_two_text_encoders` and `has_three_text_encoders` cannot be True.")


Review comment (performance):
It seems like we mostly use a single scheduler for testing with a few exceptions (CogVideoX) 

The LoRA test where scheduler type matters would be LCM LoRAs, and in those tests it looks like LCM scheduler is set directly?
https://github.com/huggingface/diffusers/blob/4067d6c4b64f2b606f9806d4a8b15d5fd5cbea1e/tests/lora/test_lora_layers_sdxl.py#L227

So why not refactor this entirely to just work with a single scheduler, and we can use the pipeline class name as the cache key? 

Code after:
    text_encoder_3_cls, text_encoder_3_id, text_encoder_3_subfolder = None, None, ""
    tokenizer_cls, tokenizer_id, tokenizer_subfolder = None, None, ""
    tokenizer_2_cls, tokenizer_2_id, tokenizer_2_subfolder = None, None, ""
    tokenizer_3_cls, tokenizer_3_id, tokenizer_3_subfolder = None, None, ""

    unet_kwargs = None
    transformer_cls = None
    transformer_kwargs = None
    vae_cls = AutoencoderKL
    vae_kwargs = None

    text_encoder_target_modules = ["q_proj", "k_proj", "v_proj", "out_proj"]
    denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0"]

    cached_non_lora_output = None

    def get_base_pipe_output(self):
        if self.cached_non_lora_output is None:
            self.cached_non_lora_output = self._compute_baseline_output()
        return self.cached_non_lora_output

    def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=None):
        if self.unet_kwargs and self.transformer_kwargs:
            raise ValueError("Both `unet_kwargs` and `transformer_kwargs` cannot be specified.")
        if self.has_two_text_encoders and self.has_three_text_encoders:
            raise ValueError("Both `has_two_text_encoders` and `has_three_text_encoders` cannot be True.")

        scheduler_cls = scheduler_cls if scheduler_cls is not None else self.scheduler_cls
        rank = 4
        lora_alpha = rank if lora_alpha is None else lora_alpha

        torch.manual_seed(0)
        if self.unet_kwargs is not None:
            unet = UNet2DConditionModel(**self.unet_kwargs)
        else:
            transformer = self.transformer_cls(**self.transformer_kwargs)

        scheduler = scheduler_cls(**self.scheduler_kwargs)

        torch.manual_seed(0)
        vae = self.vae_cls(**self.vae_kwargs)

        text_encoder = self.text_encoder_cls.from_pretrained(
            self.text_encoder_id, subfolder=self.text_encoder_subfolder
        )
        tokenizer = self.tokenizer_cls.from_pretrained(self.tokenizer_id, subfolder=self.tokenizer_subfolder)

        if self.text_encoder_2_cls is not None:
            text_encoder_2 = self.text_encoder_2_cls.from_pretrained(
                self.text_encoder_2_id, subfolder=self.text_encoder_2_subfolder
            )