Added initial support for f-lite model

extensions_built_in/diffusion_models/f_light/f_light.py

@@ -0,0 +1,295 @@
+import os
+from typing import TYPE_CHECKING
+
+import torch
+import yaml
+from toolkit.config_modules import GenerateImageConfig, ModelConfig
+from PIL import Image
+from toolkit.models.base_model import BaseModel
+from toolkit.basic import flush
+from diffusers import AutoencoderKL
+from toolkit.prompt_utils import PromptEmbeds
+from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
+from toolkit.dequantize import patch_dequantization_on_save
+from toolkit.accelerator import unwrap_model
+from optimum.quanto import freeze, QTensor
+from toolkit.util.quantize import quantize, get_qtype
+from transformers import T5TokenizerFast, T5EncoderModel
+from .src import FLitePipeline, DiT
+
+if TYPE_CHECKING:
+    from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
+
+scheduler_config = {
+    "base_image_seq_len": 256,
+    "base_shift": 0.5,
+    "max_image_seq_len": 4096,
+    "max_shift": 1.15,
+    "num_train_timesteps": 1000,
+    "shift": 3.0,
+    "use_dynamic_shifting": True
+}
+
+
+class FLiteModel(BaseModel):
+    arch = "f-lite"
+
+    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 = ['DiT']
+
+    # static method to get the noise scheduler
+    @staticmethod
+    def get_train_scheduler():
+        return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)
+    
+    def get_bucket_divisibility(self):
+        # return the bucket divisibility for the model
+        return 16
+
+    def load_model(self):
+        dtype = self.torch_dtype
+        
+        # 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
+
+        self.print_and_status_update("Loading transformer")
+
+        transformer = DiT.from_pretrained(
+            model_path,
+            subfolder="dit_model",
+            torch_dtype=dtype,
+        )
+        
+        transformer.to(self.quantize_device, dtype=dtype)
+
+        if self.model_config.quantize:
+            # patch the state dict method
+            patch_dequantization_on_save(transformer)
+            quantization_type = get_qtype(self.model_config.qtype)
+            self.print_and_status_update("Quantizing transformer")
+            quantize(transformer, weights=quantization_type,
+                     **self.model_config.quantize_kwargs)
+            freeze(transformer)
+            transformer.to(self.device_torch)
+        else:
+            transformer.to(self.device_torch, dtype=dtype)
+
+        flush()
+
+        self.print_and_status_update("Loading T5")
+        tokenizer = T5TokenizerFast.from_pretrained(
+            extras_path, subfolder="tokenizer", torch_dtype=dtype
+        )
+        text_encoder = T5EncoderModel.from_pretrained(
+            extras_path, subfolder="text_encoder", torch_dtype=dtype
+        )
+        text_encoder.to(self.device_torch, dtype=dtype)
+        flush()
+
+        if self.model_config.quantize_te:
+            self.print_and_status_update("Quantizing T5")
+            quantize(text_encoder, weights=get_qtype(
+                self.model_config.qtype))
+            freeze(text_encoder)
+            flush()
+
+        self.noise_scheduler = FLiteModel.get_train_scheduler()
+        
+        self.print_and_status_update("Loading VAE")
+        vae = AutoencoderKL.from_pretrained(
+            extras_path,
+            subfolder="vae",
+            torch_dtype=dtype
+        )
+        vae = vae.to(self.device_torch, dtype=dtype)
+
+        self.print_and_status_update("Making pipe")
+
+        pipe: FLitePipeline = FLitePipeline(
+            text_encoder=None,
+            tokenizer=tokenizer,
+            vae=vae,
+            dit_model=None,
+        )
+        # for quantization, it works best to do these after making the pipe
+        pipe.text_encoder = text_encoder
+        pipe.dit_model = transformer
+        pipe.transformer = transformer
+        pipe.scheduler = self.noise_scheduler,
+
+        self.print_and_status_update("Preparing Model")
+
+        text_encoder = [pipe.text_encoder]
+        tokenizer = [pipe.tokenizer]
+
+        pipe.transformer = pipe.transformer.to(self.device_torch)
+
+        flush()
+        # just to make sure everything is on the right device and dtype
+        text_encoder[0].to(self.device_torch)
+        text_encoder[0].requires_grad_(False)
+        text_encoder[0].eval()
+        pipe.transformer = pipe.transformer.to(self.device_torch)
+        flush()
+
+        # save it to the model class
+        self.vae = vae
+        self.text_encoder = text_encoder  # list of text encoders
+        self.tokenizer = tokenizer  # list of tokenizers
+        self.model = pipe.transformer
+        self.pipeline = pipe
+        self.print_and_status_update("Model Loaded")
+
+    def get_generation_pipeline(self):
+        scheduler = FLiteModel.get_train_scheduler()
+        # it has built in scheduler. Basically euler flowmatching
+        pipeline = FLitePipeline(
+            text_encoder=unwrap_model(self.text_encoder[0]),
+            tokenizer=self.tokenizer[0],
+            vae=unwrap_model(self.vae),
+            dit_model=unwrap_model(self.transformer)
+        )
+        pipeline.transformer = pipeline.dit_model
+        pipeline.scheduler = scheduler
+
+        return pipeline
+
+    def generate_single_image(
+        self,
+        pipeline: FLitePipeline,
+        gen_config: GenerateImageConfig,
+        conditional_embeds: PromptEmbeds,
+        unconditional_embeds: PromptEmbeds,
+        generator: torch.Generator,
+        extra: dict,
+    ):
+
+        extra['negative_prompt_embeds'] = unconditional_embeds.text_embeds
+        
+        img = pipeline(
+            prompt_embeds=conditional_embeds.text_embeds,
+            negative_prompt_embeds=unconditional_embeds.text_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,
+        ).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
+    ):
+        cast_dtype = self.unet.dtype
+
+        noise_pred = self.unet(
+            latent_model_input.to(
+                self.device_torch, cast_dtype
+            ),
+            text_embeddings.text_embeds.to(
+                self.device_torch, cast_dtype
+            ),
+            timestep / 1000,
+        )
+
+        if isinstance(noise_pred, QTensor):
+            noise_pred = noise_pred.dequantize()
+        
+        return noise_pred
+    
+    def get_prompt_embeds(self, prompt: str) -> PromptEmbeds:
+        if isinstance(prompt, str):
+            prompts = [prompt]
+        else:
+            prompts = prompt
+        if self.pipeline.text_encoder.device != self.device_torch:
+            self.pipeline.text_encoder.to(self.device_torch)
+
+        prompt_embeds, negative_embeds = self.pipeline.encode_prompt(
+            prompt=prompts, 
+            negative_prompt=None, 
+            device=self.text_encoder[0].device, 
+            dtype=self.torch_dtype,
+        )
+        
+        pe = PromptEmbeds(prompt_embeds)
+        
+        return pe
+    
+    def get_model_has_grad(self):
+        # return from a weight if it has grad
+        return False
+
+    def get_te_has_grad(self):
+        # return from a weight if it has grad
+        return False
+    
+    def save_model(self, output_path, meta, save_dtype):
+        # only save the unet
+        transformer: DiT = unwrap_model(self.model)
+        # diffusers
+        # only save the unet
+        transformer: DiT = unwrap_model(self.transformer)
+        transformer.save_pretrained(
+            save_directory=os.path.join(output_path, 'dit_model'),
+            safe_serialization=True,
+        )
+        # save out meta config
+        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()
+        return (batch.latents - noise).detach()
+    
+    def convert_lora_weights_before_save(self, state_dict):
+        # currently starte with transformer. but needs to start with diffusion_model. for comfyui
+        new_sd = {}
+        for key, value in state_dict.items():
+            new_key = key.replace("transformer.", "diffusion_model.")
+            new_sd[new_key] = value
+        return new_sd
+
+    def convert_lora_weights_before_load(self, state_dict):
+        # saved as diffusion_model. but needs to be transformer. for ai-toolkit
+        new_sd = {}
+        for key, value in state_dict.items():
+            new_key = key.replace("diffusion_model.", "transformer.")
+            new_sd[new_key] = value
+        return new_sd
+    
+    def get_base_model_version(self):
+        return "f-lite"
+    
+    def get_stepped_pred(self, pred, noise):
+        # just used for DFE support
+        latents = pred + noise
+        return latents