Method load_diffusion_model

models/hidream.py:71–124

(self)

        return getattr(self.diffusers_pipeline, name)

    def load_diffusion_model(self):
        dtype = self.model_config['dtype']
        transformer_dtype = self.model_config.get('transformer_dtype', dtype)

        llama3_path = self.model_config['llama3_path']
        if self.model_config.get('llama3_4bit', False):
            quantization_config = transformers.BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type='nf4',
                bnb_4bit_compute_dtype=dtype,
            )
        else:
            quantization_config = None
        text_encoder_4 = LlamaForCausalLM.from_pretrained(
            llama3_path,
            output_hidden_states=True,
            quantization_config=quantization_config,
            torch_dtype=dtype,
        )
        for p in text_encoder_4.parameters():
            p.requires_grad_(False)
            p.data = p.data.to('cpu')
        empty_cuda_cache()
        self.diffusers_pipeline.text_encoder_4 = text_encoder_4

        if transformer_dtype == 'nf4':
            quantization_config = diffusers.BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type='nf4',
                bnb_4bit_compute_dtype=dtype,
                llm_int8_skip_modules=KEEP_IN_HIGH_PRECISION,
            )
        else:
            quantization_config = None
        self.diffusers_pipeline.transformer = HiDreamImageTransformer2DModel.from_pretrained(
            self.model_config['diffusers_path'],
            subfolder='transformer',
            torch_dtype=dtype,
            quantization_config=quantization_config,
        )
        if transformer_dtype != 'nf4':
            for name, p in self.transformer.named_parameters():
                if not (any(x in name for x in KEEP_IN_HIGH_PRECISION)):
                    p.data = p.data.to(transformer_dtype)

        self.transformer.train()
        for name, p in self.transformer.named_parameters():
            p.original_name = name

        # Critically important! Official code saves MoE aux losses in global state if alpha > 0. Without special handling of
        # this, it causes massive memory leak during backward pass and immediately OOMs you.
        for module in self.transformer.modules():
            if isinstance(module, MoEGate):
                module.alpha = 0

    def get_vae(self):
        return self.vae