Class StableDiffusionXL

stable_diffusion/stable_diffusion/__init__.py:172–306

class StableDiffusionXL(StableDiffusion):
    def __init__(self, model: str = _DEFAULT_MODEL, float16: bool = False):
        super().__init__(model, float16)

        self.sampler = SimpleEulerAncestralSampler(self.diffusion_config)

        self.text_encoder_1 = self.text_encoder
        self.tokenizer_1 = self.tokenizer
        del self.tokenizer, self.text_encoder

        self.text_encoder_2 = load_text_encoder(
            model,
            float16,
            model_key="text_encoder_2",
        )
        self.tokenizer_2 = load_tokenizer(
            model,
            merges_key="tokenizer_2_merges",
            vocab_key="tokenizer_2_vocab",
        )

    def ensure_models_are_loaded(self):
        mx.eval(self.unet.parameters())
        mx.eval(self.text_encoder_1.parameters())
        mx.eval(self.text_encoder_2.parameters())
        mx.eval(self.autoencoder.parameters())

    def _get_text_conditioning(
        self,
        text: str,
        n_images: int = 1,
        cfg_weight: float = 7.5,
        negative_text: str = "",
    ):
        tokens_1 = self._tokenize(
            self.tokenizer_1,
            text,
            (negative_text if cfg_weight > 1 else None),
        )
        tokens_2 = self._tokenize(
            self.tokenizer_2,
            text,
            (negative_text if cfg_weight > 1 else None),
        )

        conditioning_1 = self.text_encoder_1(tokens_1)
        conditioning_2 = self.text_encoder_2(tokens_2)
        conditioning = mx.concatenate(
            [conditioning_1.hidden_states[-2], conditioning_2.hidden_states[-2]],
            axis=-1,
        )
        pooled_conditioning = conditioning_2.pooled_output

        if n_images > 1:
            conditioning = mx.repeat(conditioning, n_images, axis=0)
            pooled_conditioning = mx.repeat(pooled_conditioning, n_images, axis=0)

        return conditioning, pooled_conditioning