""" Flux2KleinBackend.py FLUX.2 Klein 9B KV Nunchaku + NVFP4 Qwen3 Text Encoder Backend. Optimized for 16GB GPUs (Blackwell sm_120 / sm_110): - Distilled 4-step inference with KV-cache for multi-reference editing. - Sequential GPU stage execution: Stage 1: Qwen3 NVFP4 Text Encoder on GPU -> offload to CPU. Stage 2: AutoencoderKLFlux2 VAE on GPU (if image conditioning) -> offload to CPU. Stage 3: NunchakuFlux2Transformer2DModel (FP4/INT4 SVDQuant) on GPU -> offload to CPU. Stage 4: AutoencoderKLFlux2 VAE on GPU (decode) -> offload to CPU. - All models reside on CPU when idle, leaving the GPU with 0% VRAM allocation for other concurrent workloads. """ from __future__ import annotations import gc import os import time from typing import Any, List, Optional, Union import numpy as np import PIL.Image import torch from diffusers import AutoencoderKLFlux2, FlowMatchEulerDiscreteScheduler from diffusers.pipelines.flux2.pipeline_flux2_klein_kv import ( Flux2KleinKVPipeline, compute_empirical_mu, retrieve_timesteps, ) from diffusers.pipelines.flux2.pipeline_output import Flux2PipelineOutput from nunchaku import NunchakuFlux2Transformer2DModel from nunchaku.utils import get_precision from transformers import AutoTokenizer from NVFP4TextEncoder import load_nvfp4_text_encoder def _flush(): gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() class SequentialFlux2KleinPipeline: """Wrapper that runs Flux2KleinKVPipeline stages sequentially on CUDA, offloading each component to CPU immediately after its stage finishes. """ def __init__( self, scheduler: FlowMatchEulerDiscreteScheduler, vae: AutoencoderKLFlux2, text_encoder: torch.nn.Module, tokenizer: AutoTokenizer, transformer: NunchakuFlux2Transformer2DModel, device: str | torch.device = "cuda", ): self.scheduler = scheduler self.vae = vae self.text_encoder = text_encoder self.tokenizer = tokenizer self.transformer = transformer self.device = torch.device(device) # Standard diffusers pipeline attributes for compatibility self.inner_pipe = Flux2KleinKVPipeline( scheduler=scheduler, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, transformer=transformer, is_distilled=True, ) self.image_processor = self.inner_pipe.image_processor self.vae_scale_factor = self.inner_pipe.vae_scale_factor # Expose components dictionary for server sleep / memory introspection self.components = { "text_encoder": self.text_encoder, "transformer": self.transformer, "vae": self.vae, "scheduler": self.scheduler, "tokenizer": self.tokenizer, } # Keep everything on CPU at initialization self.to("cpu") def to(self, device: str | torch.device): """Move components to specified device (typically 'cpu').""" target = str(device) if target == "cpu": for mod in [self.text_encoder, self.transformer, self.vae]: if isinstance(mod, torch.nn.Module): try: mod.to("cpu") except Exception: pass _flush() return self @torch.no_grad() def __call__( self, prompt: str | list[str] = "", image: Union[PIL.Image.Image, List[PIL.Image.Image], None] = None, height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 4, guidance_scale: float = 1.0, negative_prompt: Optional[str] = None, num_images_per_prompt: int = 1, generator: Optional[torch.Generator] = None, seed: Optional[int] = None, **kwargs, ) -> Flux2PipelineOutput: device = self.device if generator is None: generator = torch.Generator(device=device) if seed is not None: generator.manual_seed(seed) t_start = time.time() # ------------------------------------------------------------------ # STAGE 1: Text Encoding (Qwen3 NVFP4 Text Encoder) # ------------------------------------------------------------------ prompt_embeds = None text_ids = None try: print("[Flux2Klein] Stage 1: Moving Text Encoder to GPU...") self.text_encoder.to(device) _flush() prompt_embeds, text_ids = self.inner_pipe.encode_prompt( prompt=prompt, device=device, num_images_per_prompt=num_images_per_prompt, max_sequence_length=512, ) # Store CPU copies until Stage 3 to minimize VRAM prompt_embeds = prompt_embeds.to("cpu") text_ids = text_ids.to("cpu") finally: self.text_encoder.to("cpu") _flush() print(f"[Flux2Klein] Stage 1 complete. Allocated VRAM: {torch.cuda.memory_allocated() / 1e9:.3f} GB") # ------------------------------------------------------------------ # STAGE 2: VAE Conditioning (if reference images provided) # ------------------------------------------------------------------ condition_images = None if image is not None: if not isinstance(image, list): image = [image] condition_images = [] for img in image: self.image_processor.check_image_input(img) img_w, img_h = img.size if img_w * img_h > 1024 * 1024: img = self.image_processor._resize_to_target_area(img, 1024 * 1024) img_w, img_h = img.size multiple_of = self.vae_scale_factor * 2 img_w = (img_w // multiple_of) * multiple_of img_h = (img_h // multiple_of) * multiple_of img_pre = self.image_processor.preprocess(img, height=img_h, width=img_w, resize_mode="crop") condition_images.append(img_pre) height = height or img_h width = width or img_w height = height or 1024 width = width or 1024 image_latents = None image_latent_ids = None if condition_images is not None: try: print(f"[Flux2Klein] Stage 2: Moving VAE to GPU to encode {len(condition_images)} reference image(s)...") self.vae.to(device) _flush() image_latents, image_latent_ids = self.inner_pipe.prepare_image_latents( images=condition_images, batch_size=num_images_per_prompt, generator=generator, device=device, dtype=self.vae.dtype, ) image_latents = image_latents.to("cpu") image_latent_ids = image_latent_ids.to("cpu") finally: self.vae.to("cpu") _flush() print(f"[Flux2Klein] Stage 2 complete. Allocated VRAM: {torch.cuda.memory_allocated() / 1e9:.3f} GB") # ------------------------------------------------------------------ # STAGE 3: Transformer Denoising Loop (Nunchaku FP4/INT4 KV-Cache) # ------------------------------------------------------------------ num_channels_latents = self.transformer.config.in_channels // 4 latents, latent_ids = self.inner_pipe.prepare_latents( batch_size=num_images_per_prompt, num_latents_channels=num_channels_latents, height=height, width=width, dtype=self.transformer.dtype, device=device, generator=generator, ) sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps) image_seq_len = latents.shape[1] mu = compute_empirical_mu(image_seq_len=image_seq_len, num_steps=num_inference_steps) timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu, ) self.scheduler.set_begin_index(0) kv_cache = None try: print("[Flux2Klein] Stage 3: Moving Nunchaku Transformer to GPU...") self.transformer.to(device) _flush() prompt_embeds = prompt_embeds.to(device) text_ids = text_ids.to(device) latents = latents.to(device) latent_ids = latent_ids.to(device) if image_latents is not None: image_latents = image_latents.to(device) image_latent_ids = image_latent_ids.to(device) for i, t in enumerate(timesteps): t_step = time.time() timestep = t.expand(latents.shape[0]).to(latents.dtype) if i == 0 and image_latents is not None: # Step 0: Extract KV cache for reference tokens latent_model_input = torch.cat([image_latents, latents], dim=1).to(self.transformer.dtype) latent_image_ids = torch.cat([image_latent_ids, latent_ids], dim=1) noise_pred, kv_cache = self.transformer( hidden_states=latent_model_input, timestep=timestep / 1000, guidance=None, encoder_hidden_states=prompt_embeds, txt_ids=text_ids, img_ids=latent_image_ids, joint_attention_kwargs=None, return_dict=False, kv_cache_mode="extract", num_ref_tokens=image_latents.shape[1], ) elif kv_cache is not None: # Steps 1+: Reuse cached reference KV noise_pred = self.transformer( hidden_states=latents.to(self.transformer.dtype), timestep=timestep / 1000, guidance=None, encoder_hidden_states=prompt_embeds, txt_ids=text_ids, img_ids=latent_ids, joint_attention_kwargs=None, return_dict=False, kv_cache=kv_cache, kv_cache_mode="cached", )[0] else: # Standard text-to-image forward noise_pred = self.transformer( hidden_states=latents.to(self.transformer.dtype), timestep=timestep / 1000, guidance=None, encoder_hidden_states=prompt_embeds, txt_ids=text_ids, img_ids=latent_ids, joint_attention_kwargs=None, return_dict=False, )[0] latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] print( f"[Flux2Klein] Step {i+1}/{len(timesteps)} in {time.time() - t_step:.2f}s | " f"Peak VRAM: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB" ) if "noise_pred" in locals(): del noise_pred finally: if kv_cache is not None: kv_cache.clear() self.transformer.to("cpu") del prompt_embeds, text_ids, image_latents, image_latent_ids _flush() print(f"[Flux2Klein] Stage 3 complete. Allocated VRAM: {torch.cuda.memory_allocated() / 1e9:.3f} GB") # ------------------------------------------------------------------ # STAGE 4: VAE Decoding # ------------------------------------------------------------------ try: print("[Flux2Klein] Stage 4: Moving VAE to GPU to decode latents...") self.vae.to(device) _flush() latents = self.inner_pipe._unpack_latents_with_ids(latents, latent_ids) latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to( latents.device, latents.dtype ) latents = latents * latents_bn_std + latents_bn_mean latents = self.inner_pipe._unpatchify_latents(latents) decoded = self.vae.decode(latents, return_dict=False)[0] out_images = self.image_processor.postprocess(decoded, output_type="pil") finally: self.vae.to("cpu") del latents, latent_ids _flush() print(f"[Flux2Klein] Stage 4 complete. Allocated VRAM: {torch.cuda.memory_allocated() / 1e9:.3f} GB") print(f"[Flux2Klein] Total generation finished in {time.time() - t_start:.2f}s!") return Flux2PipelineOutput(images=out_images) class Flux2KleinBackend: """Backend factory for FLUX.2 Klein 9B KV Nunchaku.""" def __init__( self, model_id: str, optimized_model_path: Optional[str] = None, nvfp4_text_encoder_path: Optional[str] = None, ): self.model_id = model_id self.optimized_model_path = optimized_model_path self.nvfp4_text_encoder_path = nvfp4_text_encoder_path self.pipeline = None def load(self): print(f"[Flux2KleinBackend] Initializing backend from base model: {self.model_id}...") # 1. Resolve optimized transformer path opt_path = self.optimized_model_path if not opt_path: precision = get_precision() opt_path = os.path.join( self.model_id, f"svdq-{precision}_r32-FLUX.2-klein-9b-kv-Nunchaku.safetensors", ) if not os.path.exists(opt_path): # Fallback to fp4 if not found directly opt_path = os.path.join( self.model_id, "svdq-fp4_r32-FLUX.2-klein-9b-kv-Nunchaku.safetensors", ) print(f"[Flux2KleinBackend] Loading Nunchaku transformer from {opt_path}...") transformer = NunchakuFlux2Transformer2DModel.from_pretrained( opt_path, torch_dtype=torch.bfloat16, ) transformer.to("cpu") # 2. Resolve NVFP4 text encoder path enc_path = self.nvfp4_text_encoder_path if not enc_path: # Check if running directly from within the text encoder repository (extras/) candidate_parent = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) if os.path.exists(os.path.join(candidate_parent, "config.json")): enc_path = candidate_parent elif os.path.exists(os.path.join(self.model_id, "text_encoder")): enc_path = os.path.join(self.model_id, "text_encoder") else: enc_path = "/path/to/FLUX.2-klein-text-encoder-NVFP4" print(f"[Flux2KleinBackend] Loading NVFP4 text encoder from {enc_path}...") text_encoder = load_nvfp4_text_encoder(enc_path, device="cpu", dtype=torch.bfloat16) tokenizer = AutoTokenizer.from_pretrained(enc_path) # 3. Load VAE and Scheduler print(f"[Flux2KleinBackend] Loading VAE and Scheduler from {self.model_id}...") vae = AutoencoderKLFlux2.from_pretrained( self.model_id, subfolder="vae", torch_dtype=torch.bfloat16, ) vae.to("cpu") scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( self.model_id, subfolder="scheduler", ) # 4. Construct Sequential Pipeline self.pipeline = SequentialFlux2KleinPipeline( scheduler=scheduler, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, transformer=transformer, device="cuda", ) print("[Flux2KleinBackend] Pipeline successfully loaded and ready on CPU!") return self.pipeline, self.pipeline