Phase 1

extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py

@@ -6,21 +6,25 @@ import torch
 import yaml
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.vae.qwen_image import QwenImageVAE, QwenImageVAEHolderMixin
+from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder
 from toolkit.basic import flush
 from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
 from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
 from toolkit.accelerator import unwrap_model
-from optimum.quanto import freeze
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
+from toolkit.util.quantize import quantize_model
 from toolkit.memory_management import MemoryManager
 
-from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
+from transformers import Qwen3VLProcessor
 import torch.nn.functional as F
 
 try:
-    from diffusers import NucleusMoEImagePipeline, NucleusMoEImageTransformer2DModel, AutoencoderKLQwenImage
+    from diffusers import NucleusMoEImagePipeline, AutoencoderKLQwenImage
+    from toolkit.models.v2.diffusion_models.nucleus_image import (
+        NucleusMoEImageTransformer2DModel,
+    )
     from diffusers.models.transformers.transformer_nucleusmoe_image import SwiGLUExperts
 except ImportError:
     raise ImportError(
@@ -46,7 +50,7 @@ scheduler_config = {
 }
 
 
-class NucleusImageModel(BaseModel):
+class NucleusImageModel(QwenImageVAEHolderMixin, BaseModel):
     arch = "nucleus_image"
 
     def __init__(
@@ -81,20 +85,14 @@ class NucleusImageModel(BaseModel):
 
         self.print_and_status_update("Loading transformer")
 
-        transformer_path = model_path
-        transformer_subfolder = "transformer"
-        if os.path.exists(transformer_path):
-            transformer_subfolder = None
-            transformer_path = os.path.join(transformer_path, "transformer")
+        if os.path.exists(model_path):
             # check if the path is a full checkpoint.
             te_folder_path = os.path.join(model_path, "text_encoder")
             # if we have the te, this folder is a full checkpoint, use it as the base
             if os.path.exists(te_folder_path):
                 base_model_path = model_path
 
-        transformer = NucleusMoEImageTransformer2DModel.from_pretrained(
-            transformer_path, subfolder=transformer_subfolder, torch_dtype=dtype
-        )
+        transformer = NucleusMoEImageTransformer2DModel.load_model(model_path, dtype=dtype)
         
         # handle versions of pytorch that don't have grouped mm, by disabling it in the SwiGLUExperts
         if not hasattr(torch.nn.functional, "grouped_mm"):
@@ -129,33 +127,13 @@ class NucleusImageModel(BaseModel):
         tokenizer = Qwen3VLProcessor.from_pretrained(
             base_model_path, subfolder="processor", torch_dtype=dtype
         )
-        text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
-            base_model_path, subfolder="text_encoder", torch_dtype=dtype
-        )
-
-        if (
-            self.model_config.layer_offloading
-            and self.model_config.layer_offloading_text_encoder_percent > 0
-        ):
-            MemoryManager.attach(
-                text_encoder,
-                self.device_torch,
-                offload_percent=self.model_config.layer_offloading_text_encoder_percent,
-            )
-
-        text_encoder.to(self.device_torch, dtype=dtype)
-        flush()
-
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing Text Encoder")
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(text_encoder)
-            flush()
+        text_encoder = Qwen3VLTextEncoder.load_model(base_model_path, dtype=dtype)
+        self.prepare_text_encoder(text_encoder, dtype=dtype)
 
         self.print_and_status_update("Loading VAE")
-        vae = AutoencoderKLQwenImage.from_pretrained(
-            base_model_path, subfolder="vae", torch_dtype=dtype
-        ).to(self.device_torch, dtype=dtype)
+        vae = QwenImageVAE.load_model(
+            base_model_path, dtype=dtype, device=self.device_torch
+        )
 
         self.noise_scheduler = NucleusImageModel.get_train_scheduler()
 
@@ -214,41 +192,6 @@ class NucleusImageModel(BaseModel):
 
         return pipeline
 
-    def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None):
-        if device is None:
-            device = self.vae_device_torch
-        if dtype is None:
-            dtype = self.vae_torch_dtype
-
-        # Move to vae to device if on cpu
-        if self.vae.device == torch.device("cpu"):
-            self.vae.to(device)
-        self.vae.eval()
-        self.vae.requires_grad_(False)
-        # move to device and dtype
-        image_list = [image.to(device, dtype=dtype) for image in image_list]
-        images = torch.stack(image_list).to(device, dtype=dtype)
-        # it uses wan vae, so add dim for frame count
-
-        images = images.unsqueeze(2)
-        latents = self.vae.encode(images).latent_dist.sample()
-
-        latents_mean = (
-            torch.tensor(self.vae.config.latents_mean)
-            .view(1, self.vae.config.z_dim, 1, 1, 1)
-            .to(latents.device, latents.dtype)
-        )
-        latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(
-            1, self.vae.config.z_dim, 1, 1, 1
-        ).to(latents.device, latents.dtype)
-
-        latents = (latents - latents_mean) * latents_std
-        latents = latents.to(device, dtype=dtype)
-
-        latents = latents.squeeze(2)  # remove the frame count dimension
-
-        return latents
-
     def generate_single_image(
         self,
         pipeline: NucleusMoEImagePipeline,