import os from typing import TYPE_CHECKING, List, Optional import einops import torch import torchvision import yaml from toolkit import train_tools from toolkit.config_modules import GenerateImageConfig, ModelConfig from PIL import Image from toolkit.models.base_model import BaseModel from toolkit.models.v2.text_encoders.t5 import T5TextEncoder from toolkit.models.v2.text_encoders.llama import LlamaTextEncoder from toolkit.models.v2.text_encoders.clip import CLIPTextEncoderWithProjection from toolkit.models.v2.vae.autoencoder_kl import KLVAE from diffusers import TorchAoConfig from toolkit.basic import flush from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler from toolkit.models.flux import add_model_gpu_splitter_to_flux, bypass_flux_guidance, restore_flux_guidance from toolkit.dequantize import patch_dequantization_on_save from toolkit.accelerator import get_accelerator, unwrap_model from optimum.quanto import freeze, QTensor from toolkit.util.mask import generate_random_mask, random_dialate_mask from .src.pipelines.hidream_image.pipeline_hidream_image import HiDreamImagePipeline from .src.models.transformers.transformer_hidream_image import HiDreamImageTransformer2DModel from .src.schedulers.fm_solvers_unipc import FlowUniPCMultistepScheduler from transformers import LlamaForCausalLM, PreTrainedTokenizerFast from einops import rearrange, repeat import random import torch.nn.functional as F from tqdm import tqdm if TYPE_CHECKING: from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO scheduler_config = { "num_train_timesteps": 1000, "shift": 3.0 } # LLAMA_MODEL_NAME = "meta-llama/Meta-Llama-3.1-8B-Instruct" LLAMA_MODEL_PATH = "unsloth/Meta-Llama-3.1-8B-Instruct" BASE_MODEL_PATH = "HiDream-ai/HiDream-I1-Full" class HidreamModel(BaseModel): arch = "hidream" hidream_transformer_class = HiDreamImageTransformer2DModel hidream_pipeline_class = HiDreamImagePipeline def __init__( self, device, model_config: ModelConfig, dtype='bf16', custom_pipeline=None, noise_scheduler=None, **kwargs ): super().__init__( device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs ) self.is_flow_matching = True self.is_transformer = True self.target_lora_modules = ['HiDreamImageTransformer2DModel'] # static method to get the noise scheduler @staticmethod def get_train_scheduler(): return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config) def get_bucket_divisibility(self): return 16 def load_model(self): dtype = self.torch_dtype # HiDream-ai/HiDream-I1-Full self.print_and_status_update("Loading HiDream model") # will be updated if we detect a existing checkpoint in training folder model_path = self.model_config.name_or_path extras_path = self.model_config.extras_name_or_path llama_model_path = self.model_config.model_kwargs.get('llama_model_path', LLAMA_MODEL_PATH) scheduler = HidreamModel.get_train_scheduler() self.print_and_status_update("Loading llama 8b model") tokenizer_4 = PreTrainedTokenizerFast.from_pretrained( llama_model_path, use_fast=False ) # load + quantize + offload + placement, all driven by model_config text_encoder_4 = LlamaTextEncoder.load( llama_model_path, subfolder="", output_hidden_states=True, output_attentions=True, **self.component_load_kwargs("te"), ) flush() self.print_and_status_update("Loading transformer") transformer = self.hidream_transformer_class.load( model_path, **self.component_load_kwargs("transformer") ) flush() self.print_and_status_update("Loading vae") vae = KLVAE.load_model(extras_path, dtype=torch.bfloat16).to( self.device_torch, dtype=dtype ) self.print_and_status_update("Loading clip encoders") text_encoder = CLIPTextEncoderWithProjection.load_model( extras_path, dtype=torch.bfloat16 ).to(self.device_torch, dtype=dtype) tokenizer = CLIPTextEncoderWithProjection.load_tokenizer( extras_path, use_fast=False ) text_encoder_2 = CLIPTextEncoderWithProjection.load_model( extras_path, dtype=torch.bfloat16, subfolder="text_encoder_2" ).to(self.device_torch, dtype=dtype) tokenizer_2 = CLIPTextEncoderWithProjection.load_tokenizer( extras_path, subfolder="tokenizer_2", use_fast=False ) flush() self.print_and_status_update("Loading T5 encoders") # load + quantize + offload + placement, all driven by model_config text_encoder_3 = T5TextEncoder.load( extras_path, subfolder="text_encoder_3", **self.component_load_kwargs("te") ) flush() tokenizer_3 = T5TextEncoder.load_tokenizer( extras_path, subfolder="tokenizer_3", use_fast=False ) flush() if self.low_vram: self.print_and_status_update("Moving everything to device") # move it all back transformer.to(self.device_torch, dtype=dtype) vae.to(self.device_torch, dtype=dtype) text_encoder.to(self.device_torch, dtype=dtype) text_encoder_2.to(self.device_torch, dtype=dtype) text_encoder_4.to(self.device_torch, dtype=dtype) text_encoder_3.to(self.device_torch, dtype=dtype) # set to eval mode # transformer.eval() vae.eval() text_encoder.eval() text_encoder_2.eval() text_encoder_4.eval() text_encoder_3.eval() pipe = self.hidream_pipeline_class( scheduler=scheduler, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, text_encoder_2=text_encoder_2, tokenizer_2=tokenizer_2, text_encoder_3=text_encoder_3, tokenizer_3=tokenizer_3, text_encoder_4=text_encoder_4, tokenizer_4=tokenizer_4, transformer=transformer, ) flush() text_encoder_list = [text_encoder, text_encoder_2, text_encoder_3, text_encoder_4] tokenizer_list = [tokenizer, tokenizer_2, tokenizer_3, tokenizer_4] for te in text_encoder_list: # set the dtype te.to(self.device_torch, dtype=dtype) # freeze the model freeze(te) # set to eval mode te.eval() # set the requires grad to false te.requires_grad_(False) flush() # save it to the model class self.vae = vae self.text_encoder = text_encoder_list # list of text encoders self.tokenizer = tokenizer_list # list of tokenizers self.model = pipe.transformer self.pipeline = pipe self.print_and_status_update("Model Loaded") def get_generation_pipeline(self): scheduler = FlowUniPCMultistepScheduler( num_train_timesteps=1000, shift=3.0, use_dynamic_shifting=False ) pipeline: HiDreamImagePipeline = HiDreamImagePipeline( scheduler=scheduler, vae=self.vae, text_encoder=self.text_encoder[0], tokenizer=self.tokenizer[0], text_encoder_2=self.text_encoder[1], tokenizer_2=self.tokenizer[1], text_encoder_3=self.text_encoder[2], tokenizer_3=self.tokenizer[2], text_encoder_4=self.text_encoder[3], tokenizer_4=self.tokenizer[3], transformer=unwrap_model(self.model), aggressive_unloading=self.low_vram ) pipeline = pipeline.to(self.device_torch) return pipeline def generate_single_image( self, pipeline: HiDreamImagePipeline, gen_config: GenerateImageConfig, conditional_embeds: PromptEmbeds, unconditional_embeds: PromptEmbeds, generator: torch.Generator, extra: dict, ): img = pipeline( prompt_embeds=conditional_embeds.text_embeds, pooled_prompt_embeds=conditional_embeds.pooled_embeds, negative_prompt_embeds=unconditional_embeds.text_embeds, negative_pooled_prompt_embeds=unconditional_embeds.pooled_embeds, height=gen_config.height, width=gen_config.width, num_inference_steps=gen_config.num_inference_steps, guidance_scale=gen_config.guidance_scale, latents=gen_config.latents, generator=generator, **extra ).images[0] return img def get_noise_prediction( self, latent_model_input: torch.Tensor, timestep: torch.Tensor, # 0 to 1000 scale text_embeddings: PromptEmbeds, **kwargs ): batch_size = latent_model_input.shape[0] with torch.no_grad(): if latent_model_input.shape[-2] != latent_model_input.shape[-1]: B, C, H, W = latent_model_input.shape pH, pW = H // self.model.config.patch_size, W // self.model.config.patch_size img_sizes = torch.tensor([pH, pW], dtype=torch.int64).reshape(-1) img_ids = torch.zeros(pH, pW, 3) img_ids[..., 1] = img_ids[..., 1] + torch.arange(pH)[:, None] img_ids[..., 2] = img_ids[..., 2] + torch.arange(pW)[None, :] img_ids = img_ids.reshape(pH * pW, -1) img_ids_pad = torch.zeros(self.transformer.max_seq, 3) img_ids_pad[:pH*pW, :] = img_ids img_sizes = img_sizes.unsqueeze(0).to(latent_model_input.device) img_sizes = torch.cat([img_sizes] * batch_size, dim=0) img_ids = img_ids_pad.unsqueeze(0).to(latent_model_input.device) img_ids = torch.cat([img_ids] * batch_size, dim=0) else: img_sizes = img_ids = None dtype = self.model.dtype device = self.device_torch # Pack the latent if latent_model_input.shape[-2] != latent_model_input.shape[-1]: B, C, H, W = latent_model_input.shape patch_size = self.transformer.config.patch_size pH, pW = H // patch_size, W // patch_size out = torch.zeros( (B, C, self.transformer.max_seq, patch_size * patch_size), dtype=latent_model_input.dtype, device=latent_model_input.device ) latent_model_input = einops.rearrange(latent_model_input, 'B C (H p1) (W p2) -> B C (H W) (p1 p2)', p1=patch_size, p2=patch_size) out[:, :, 0:pH*pW] = latent_model_input latent_model_input = out text_embeds = text_embeddings.text_embeds # run the to for the list text_embeds = [te.to(device, dtype=dtype) for te in text_embeds] noise_pred = self.transformer( hidden_states = latent_model_input, timesteps = timestep, encoder_hidden_states = text_embeds, pooled_embeds = text_embeddings.pooled_embeds.to(device, dtype=dtype), img_sizes = img_sizes, img_ids = img_ids, return_dict = False, )[0] noise_pred = -noise_pred return noise_pred def get_prompt_embeds(self, prompt: str) -> PromptEmbeds: self.text_encoder_to(self.device_torch, dtype=self.torch_dtype) max_sequence_length = 128 prompt_embeds, pooled_prompt_embeds = self.pipeline._encode_prompt( prompt = prompt, prompt_2 = prompt, prompt_3 = prompt, prompt_4 = prompt, device = self.device_torch, dtype = self.torch_dtype, num_images_per_prompt = 1, max_sequence_length = max_sequence_length, ) pe = PromptEmbeds( [prompt_embeds, pooled_prompt_embeds] ) return pe def get_model_has_grad(self): # return from a weight if it has grad return self.model.double_stream_blocks[0].block.attn1.to_q.weight.requires_grad def get_te_has_grad(self): # assume no one wants to finetune 4 text encoders. return False def save_model(self, output_path, meta, save_dtype): # only save the unet transformer: HiDreamImageTransformer2DModel = unwrap_model(self.model) transformer.save_pretrained( save_directory=os.path.join(output_path, 'transformer'), safe_serialization=True, ) meta_path = os.path.join(output_path, 'aitk_meta.yaml') with open(meta_path, 'w') as f: yaml.dump(meta, f) def get_loss_target(self, *args, **kwargs): noise = kwargs.get('noise') batch = kwargs.get('batch') return (noise - batch.latents).detach() def get_transformer_block_names(self) -> Optional[List[str]]: return ['double_stream_blocks', 'single_stream_blocks'] lora_keys_use_comfy_prefix = True def get_base_model_version(self): return "hidream_i1"