Z-Image-Turbo-Text-Encoder-Heretic-NVFP4

Duplicate from catplusplus/Z-Image-Turbo-Text-Encoder-Heretic-NVFP4

d00f17b 5 months ago

6.08 kB

import os
import torch
from diffusers import ZImagePipeline
from nunchaku.models.transformers.transformer_zimage import NunchakuZImageTransformer2DModel
from nunchaku.utils import get_gpu_memory


class ZImageTurboBackend:
    def __init__(
        self,
        model_id,
        optimized_model_path=None,
        optimized_edit_model_path=None,
        uma=False,
        nvfp4_text_encoder_path: str | None = None,
    ):
        self.model_id = model_id
        self.optimized_model_path = optimized_model_path
        self.pipeline = None
        self.uma = uma
        # Optional path to an NVFP4-pack-quantized Qwen3 text encoder. When set,
        # we load the encoder via vLLM's CompressedTensorsW4A4Fp4 (CUTLASS NVFP4
        # GEMM) instead of the bf16 text_encoder shipped inside the Z-Image
        # base repo. Cuts encoder VRAM ~4x with negligible quality loss
        # (cosine >0.999 vs the bf16 reference on Thor).
        self.nvfp4_text_encoder_path = nvfp4_text_encoder_path

    def _build_nvfp4_text_encoder(self):
        """Load the NVFP4 text encoder if requested, returns (encoder, tokenizer) or (None, None)."""
        if not self.nvfp4_text_encoder_path:
            return None, None
        print(
            f"[ZImageTurboBackend] Loading NVFP4 text encoder from {self.nvfp4_text_encoder_path} "
            "(vLLM CompressedTensorsW4A4Fp4 + CUTLASS NVFP4 GEMM)"
        )
        from NVFP4TextEncoder import load_nvfp4_text_encoder
        from transformers import AutoTokenizer

        encoder = load_nvfp4_text_encoder(
            self.nvfp4_text_encoder_path,
            device="cuda",
            dtype=torch.bfloat16,
        )
        tokenizer = AutoTokenizer.from_pretrained(self.nvfp4_text_encoder_path)
        return encoder, tokenizer

    def load(self):
        print(f"Loading ZImageTurboBackend from {self.model_id}...")
        print(f"Loading NunchakuZImageTransformer2DModel from {self.optimized_model_path}...")

        # Load transformer (optimized model)
        transformer = NunchakuZImageTransformer2DModel.from_pretrained(self.optimized_model_path)

        # If requested, build the NVFP4 text encoder before constructing the pipeline so
        # diffusers does not also load the bf16 text_encoder from disk (it would double VRAM).
        nvfp4_encoder, nvfp4_tokenizer = self._build_nvfp4_text_encoder()

        # Load pipeline
        print("Initializing ZImagePipeline...")
        pipeline_kwargs = dict(
            transformer=transformer,
            torch_dtype=torch.bfloat16,
            low_cpu_mem_usage=False,  # standard for HF example
        )
        if nvfp4_encoder is not None:
            # Pass our pre-built encoder so diffusers skips loading the bf16 subfolder.
            pipeline_kwargs["text_encoder"] = nvfp4_encoder
            if nvfp4_tokenizer is not None:
                pipeline_kwargs["tokenizer"] = nvfp4_tokenizer

        pipeline = ZImagePipeline.from_pretrained(self.model_id, **pipeline_kwargs)

        gpu_mem = get_gpu_memory()
        print(f"GPU memory available: {gpu_mem} GB")

        # Enable Flash Attention 2
        try:
            if hasattr(pipeline.transformer, "set_attention_backend"):
                pipeline.transformer.set_attention_backend("native")
                print("Enabled Native SDPA for Z-Image transformer")
            if hasattr(pipeline.vae, "set_attention_backend"):
                pipeline.vae.set_attention_backend("native")
                print("Enabled Native SDPA for Z-Image VAE")
        except Exception as e:
            print(f"Could not enable Flash Attention 2: {e}")

        if self.uma:
            print("UMA mode enabled: Loading all components to GPU and disabling offloads")
            # When using the NVFP4 encoder, it is already on CUDA and its quantised parameters
            # are not compatible with diffusers' generic .to() pathway (e.g. uint8 weight_packed).
            # We move only the diffusers-managed components (vae, transformer if not nunchaku, ...).
            if nvfp4_encoder is not None:
                # Exclude text_encoder from blanket .to('cuda'); it is already on cuda.
                excl = getattr(pipeline, "_exclude_from_cpu_offload", [])
                if "text_encoder" not in excl:
                    excl.append("text_encoder")
                    pipeline._exclude_from_cpu_offload = excl
                for name, comp in pipeline.components.items():
                    if name == "text_encoder":
                        continue
                    if isinstance(comp, torch.nn.Module):
                        try:
                            comp.to("cuda")
                        except Exception:
                            pass
            else:
                pipeline.to("cuda")
        elif gpu_mem <= 18:
            print("GPU memory <= 18GB, using sequential cpu offload for low VRAM")
            # The prompt requested sequential offloading without splitting layers for Nunchaku
            pipeline._exclude_from_cpu_offload.append("transformer")
            if nvfp4_encoder is not None:
                # NVFP4 weights live entirely on CUDA; do not let accelerate move them.
                pipeline._exclude_from_cpu_offload.append("text_encoder")
            pipeline.enable_sequential_cpu_offload()
            transformer.to("cuda")
            if nvfp4_encoder is not None:
                nvfp4_encoder.to("cuda")
        else:
            print("GPU memory > 18GB, using cpu offload")
            if nvfp4_encoder is not None:
                if not hasattr(pipeline, "_exclude_from_cpu_offload"):
                    pipeline._exclude_from_cpu_offload = []
                pipeline._exclude_from_cpu_offload.append("text_encoder")
            pipeline.enable_model_cpu_offload()
            if nvfp4_encoder is not None:
                nvfp4_encoder.to("cuda")

        self.pipeline = pipeline
        # Return twice for pipeline and edit_pipeline (though Z-Image-Turbo is T2I only)
        return self.pipeline, self.pipeline

	import os
	import torch
	from diffusers import ZImagePipeline
	from nunchaku.models.transformers.transformer_zimage import NunchakuZImageTransformer2DModel
	from nunchaku.utils import get_gpu_memory


	class ZImageTurboBackend:
	def __init__(
	self,
	model_id,
	optimized_model_path=None,
	optimized_edit_model_path=None,
	uma=False,
	nvfp4_text_encoder_path: str \| None = None,
	):
	self.model_id = model_id
	self.optimized_model_path = optimized_model_path
	self.pipeline = None
	self.uma = uma
	# Optional path to an NVFP4-pack-quantized Qwen3 text encoder. When set,
	# we load the encoder via vLLM's CompressedTensorsW4A4Fp4 (CUTLASS NVFP4
	# GEMM) instead of the bf16 text_encoder shipped inside the Z-Image
	# base repo. Cuts encoder VRAM ~4x with negligible quality loss
	# (cosine >0.999 vs the bf16 reference on Thor).
	self.nvfp4_text_encoder_path = nvfp4_text_encoder_path

	def _build_nvfp4_text_encoder(self):
	"""Load the NVFP4 text encoder if requested, returns (encoder, tokenizer) or (None, None)."""
	if not self.nvfp4_text_encoder_path:
	return None, None
	print(
	f"[ZImageTurboBackend] Loading NVFP4 text encoder from {self.nvfp4_text_encoder_path} "
	"(vLLM CompressedTensorsW4A4Fp4 + CUTLASS NVFP4 GEMM)"
	)
	from NVFP4TextEncoder import load_nvfp4_text_encoder
	from transformers import AutoTokenizer

	encoder = load_nvfp4_text_encoder(
	self.nvfp4_text_encoder_path,
	device="cuda",
	dtype=torch.bfloat16,
	)
	tokenizer = AutoTokenizer.from_pretrained(self.nvfp4_text_encoder_path)
	return encoder, tokenizer

	def load(self):
	print(f"Loading ZImageTurboBackend from {self.model_id}...")
	print(f"Loading NunchakuZImageTransformer2DModel from {self.optimized_model_path}...")

	# Load transformer (optimized model)
	transformer = NunchakuZImageTransformer2DModel.from_pretrained(self.optimized_model_path)

	# If requested, build the NVFP4 text encoder before constructing the pipeline so
	# diffusers does not also load the bf16 text_encoder from disk (it would double VRAM).
	nvfp4_encoder, nvfp4_tokenizer = self._build_nvfp4_text_encoder()

	# Load pipeline
	print("Initializing ZImagePipeline...")
	pipeline_kwargs = dict(
	transformer=transformer,
	torch_dtype=torch.bfloat16,
	low_cpu_mem_usage=False, # standard for HF example
	)
	if nvfp4_encoder is not None:
	# Pass our pre-built encoder so diffusers skips loading the bf16 subfolder.
	pipeline_kwargs["text_encoder"] = nvfp4_encoder
	if nvfp4_tokenizer is not None:
	pipeline_kwargs["tokenizer"] = nvfp4_tokenizer

	pipeline = ZImagePipeline.from_pretrained(self.model_id, **pipeline_kwargs)

	gpu_mem = get_gpu_memory()
	print(f"GPU memory available: {gpu_mem} GB")

	# Enable Flash Attention 2
	try:
	if hasattr(pipeline.transformer, "set_attention_backend"):
	pipeline.transformer.set_attention_backend("native")
	print("Enabled Native SDPA for Z-Image transformer")
	if hasattr(pipeline.vae, "set_attention_backend"):
	pipeline.vae.set_attention_backend("native")
	print("Enabled Native SDPA for Z-Image VAE")
	except Exception as e:
	print(f"Could not enable Flash Attention 2: {e}")

	if self.uma:
	print("UMA mode enabled: Loading all components to GPU and disabling offloads")
	# When using the NVFP4 encoder, it is already on CUDA and its quantised parameters
	# are not compatible with diffusers' generic .to() pathway (e.g. uint8 weight_packed).
	# We move only the diffusers-managed components (vae, transformer if not nunchaku, ...).
	if nvfp4_encoder is not None:
	# Exclude text_encoder from blanket .to('cuda'); it is already on cuda.
	excl = getattr(pipeline, "_exclude_from_cpu_offload", [])
	if "text_encoder" not in excl:
	excl.append("text_encoder")
	pipeline._exclude_from_cpu_offload = excl
	for name, comp in pipeline.components.items():
	if name == "text_encoder":
	continue
	if isinstance(comp, torch.nn.Module):
	try:
	comp.to("cuda")
	except Exception:
	pass
	else:
	pipeline.to("cuda")
	elif gpu_mem <= 18:
	print("GPU memory <= 18GB, using sequential cpu offload for low VRAM")
	# The prompt requested sequential offloading without splitting layers for Nunchaku
	pipeline._exclude_from_cpu_offload.append("transformer")
	if nvfp4_encoder is not None:
	# NVFP4 weights live entirely on CUDA; do not let accelerate move them.
	pipeline._exclude_from_cpu_offload.append("text_encoder")
	pipeline.enable_sequential_cpu_offload()
	transformer.to("cuda")
	if nvfp4_encoder is not None:
	nvfp4_encoder.to("cuda")
	else:
	print("GPU memory > 18GB, using cpu offload")
	if nvfp4_encoder is not None:
	if not hasattr(pipeline, "_exclude_from_cpu_offload"):
	pipeline._exclude_from_cpu_offload = []
	pipeline._exclude_from_cpu_offload.append("text_encoder")
	pipeline.enable_model_cpu_offload()
	if nvfp4_encoder is not None:
	nvfp4_encoder.to("cuda")

	self.pipeline = pipeline
	# Return twice for pipeline and edit_pipeline (though Z-Image-Turbo is T2I only)
	return self.pipeline, self.pipeline