L2P-open-source-final/diffsynth/pipelines/z_image_L2P.py

import torch, math
from PIL import Image
from typing import Union
from tqdm import tqdm
from einops import rearrange
import numpy as np
from typing import Union, List, Optional, Tuple

from ..diffusion import FlowMatchScheduler
from ..core import ModelConfig, gradient_checkpoint_forward
# from ..diffusion.base_pipeline import BasePipeline, PipelineUnit, ControlNetInput
from ..diffusion.base_pipeline_L2P import BasePipeline, PipelineUnit, ControlNetInput

from transformers import AutoTokenizer
from ..models.z_image_text_encoder import ZImageTextEncoder
from ..models.z_image_dit_L2P import ZImageDiT
# from ..models.flux_vae import FluxVAEEncoder, FluxVAEDecoder


class ZImagePipeline(BasePipeline):

    def __init__(self, device="cuda", torch_dtype=torch.bfloat16):
        super().__init__(
            device=device, torch_dtype=torch_dtype,
            height_division_factor=16, width_division_factor=16,
        )
        self.scheduler = FlowMatchScheduler("Z-Image")
        self.text_encoder: ZImageTextEncoder = None
        self.dit: ZImageDiT = None
        # self.vae_encoder: FluxVAEEncoder = None
        # self.vae_decoder: FluxVAEDecoder = None
        self.tokenizer: AutoTokenizer = None
        self.in_iteration_models = ("dit",)
        self.units = [
            ZImageUnit_ShapeChecker(),
            ZImageUnit_PromptEmbedder(),
            ZImageUnit_NoiseInitializer(),
            ZImageUnit_InputImageEmbedder(),
        ]
        self.model_fn = model_fn_z_image
    
    
    @staticmethod
    def from_pretrained(
        torch_dtype: torch.dtype = torch.bfloat16,
        device: Union[str, torch.device] = "cuda",
        model_configs: list[ModelConfig] = [],
        tokenizer_config: ModelConfig = ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="tokenizer/"),
        vram_limit: float = None,
    ):
        # Initialize pipeline
        pipe = ZImagePipeline(device=device, torch_dtype=torch_dtype)
        model_pool = pipe.download_and_load_models(model_configs, vram_limit)
        
        # Fetch models
        pipe.text_encoder = model_pool.fetch_model("z_image_text_encoder")
        pipe.dit = model_pool.fetch_model("z_image_dit_L2P")
        # pipe.vae_encoder = model_pool.fetch_model("flux_vae_encoder")
        # pipe.vae_decoder = model_pool.fetch_model("flux_vae_decoder")
        if tokenizer_config is not None:
            tokenizer_config.download_if_necessary()
            pipe.tokenizer = AutoTokenizer.from_pretrained(tokenizer_config.path)
        
        # VRAM Management
        pipe.vram_management_enabled = pipe.check_vram_management_state()
        return pipe
    
    
    @torch.no_grad()
    def __call__(
        self,
        # Prompt
        prompt: str,
        negative_prompt: str = "",
        cfg_scale: float = 1.0,
        # Image
        input_image: Image.Image = None,
        denoising_strength: float = 1.0,
        # Shape
        height: int = 1024,
        width: int = 1024,
        # Randomness
        seed: int = None,
        rand_device: str = "cpu",
        # Steps
        num_inference_steps: int = 8,
        # Progress bar
        progress_bar_cmd = tqdm,
    ):
        # Scheduler
        self.scheduler.set_timesteps(num_inference_steps, denoising_strength=denoising_strength)
        
        # Parameters
        inputs_posi = {
            "prompt": prompt,
        }
        inputs_nega = {
            "negative_prompt": negative_prompt,
        }
        inputs_shared = {
            "cfg_scale": cfg_scale,
            "input_image": input_image, "denoising_strength": denoising_strength,
            "height": height, "width": width,
            "seed": seed, "rand_device": rand_device,
            "num_inference_steps": num_inference_steps,
        }
        for unit in self.units:
            inputs_shared, inputs_posi, inputs_nega = self.unit_runner(unit, self, inputs_shared, inputs_posi, inputs_nega)

        # Denoise
        self.load_models_to_device(self.in_iteration_models)
        models = {name: getattr(self, name) for name in self.in_iteration_models}
        for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)):
            timestep = timestep.unsqueeze(0).to(dtype=self.torch_dtype, device=self.device)
            # print(timestep)
            noise_pred = self.cfg_guided_model_fn(
                self.model_fn, cfg_scale,
                inputs_shared, inputs_posi, inputs_nega,
                **models, timestep=timestep, progress_id=progress_id
            )
            inputs_shared["latents"] = self.step(self.scheduler, progress_id=progress_id, noise_pred=noise_pred.float(), **inputs_shared)
        
        # Decode
        # self.load_models_to_device(['vae_decoder'])
        # image = self.vae_decoder(inputs_shared["latents"])
        # image = self.vae_output_to_image(image)

        image_tensor = inputs_shared["latents"]
        image = self.pixel_output_to_image(image_tensor)
        self.load_models_to_device([])

        return image

    # def pixel_output_to_image(self, pixel_tensor):
    #         """
    #         将 Pixel Space 的 Tensor (B, C, H, W) 转换为 PIL Image
    #         假设模型输出范围为 [-1, 1]
    #         """
    #         pixel_tensor = (pixel_tensor / 2 + 0.5).clamp(0, 1)
    #         pixel_tensor = pixel_tensor.cpu().permute(0, 2, 3, 1).float().numpy()
    #         pixel_tensor = (pixel_tensor * 255).round().astype("uint8")
    #         images = [Image.fromarray(image) for image in pixel_tensor]
    #         return images[0] if len(images) == 1 else images

    def pixel_output_to_image(self, pixel_tensor, min_value=-1, max_value=1):
        """
        将 Pixel Space 的 Tensor (B, C, H, W) 转换为 PIL Image
        适配输入范围 [min_value, max_value]
        """
        range_val = max_value - min_value
        pixel_tensor = (pixel_tensor - min_value) / range_val
        
        pixel_tensor = pixel_tensor.clamp(0, 1)
        pixel_tensor = pixel_tensor.cpu().permute(0, 2, 3, 1).float().numpy()
        
        pixel_tensor = (pixel_tensor * 255).round().astype("uint8")
        
        images = [Image.fromarray(image) for image in pixel_tensor]
        return images[0] if len(images) == 1 else images
    
    

class ZImageUnit_ShapeChecker(PipelineUnit):
    def __init__(self):
        super().__init__(
            input_params=("height", "width"),
            output_params=("height", "width"),
        )

    def process(self, pipe: ZImagePipeline, height, width):
        height, width = pipe.check_resize_height_width(height, width)
        return {"height": height, "width": width}


class ZImageUnit_PromptEmbedder(PipelineUnit):
    def __init__(self):
        super().__init__(
            seperate_cfg=True,
            input_params_posi={"prompt": "prompt"},
            input_params_nega={"prompt": "negative_prompt"},
            output_params=("prompt_embeds",),
            onload_model_names=("text_encoder",)
        )
    
    def encode_prompt(
        self,
        pipe,
        prompt: Union[str, List[str]],
        device: Optional[torch.device] = None,
        max_sequence_length: int = 512,
    ) -> List[torch.FloatTensor]:
        if isinstance(prompt, str):
            prompt = [prompt]

        for i, prompt_item in enumerate(prompt):
            messages = [
                {"role": "user", "content": prompt_item},
            ]
            prompt_item = pipe.tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
                enable_thinking=True,
            )
            prompt[i] = prompt_item

        text_inputs = pipe.tokenizer(
            prompt,
            padding="max_length",
            max_length=max_sequence_length,
            truncation=True,
            return_tensors="pt",
        )

        text_input_ids = text_inputs.input_ids.to(device)
        prompt_masks = text_inputs.attention_mask.to(device).bool()

        with torch.no_grad():
            prompt_embeds = pipe.text_encoder(
                input_ids=text_input_ids,
                attention_mask=prompt_masks,
                output_hidden_states=True,
            ).hidden_states[-2]

        embeddings_list = []

        for i in range(len(prompt_embeds)):
            embeddings_list.append(prompt_embeds[i][prompt_masks[i]])

        return embeddings_list

    def process(self, pipe: ZImagePipeline, prompt):
        # ============================================================
        # Optional VRAM optimization: when --offload_text_encoder is set,
        # text_encoder is kept on CPU and only moved to GPU here for
        # encoding, then immediately moved back. Saves ~5-6GB bf16 VRAM.
        # This branch is mutually exclusive with the official
        # vram_management path (load_models_to_device).
        # Numerically equivalent: text_encoder is frozen and outputs are
        # produced under torch.no_grad (see encode_prompt above).
        # ============================================================
        manual_offload = (
            getattr(pipe, "offload_text_encoder", False)
            and not getattr(pipe, "vram_management_enabled", False)
            and pipe.text_encoder is not None
            and torch.cuda.is_available()
        )

        if manual_offload:
            # text_encoder -> GPU
            pipe.text_encoder.to(pipe.device)
            pipe.text_encoder.eval()
            try:
                prompt_embeds = self.encode_prompt(pipe, prompt, pipe.device)
            finally:
                # text_encoder -> CPU, free GPU cache immediately
                pipe.text_encoder.to("cpu")
                torch.cuda.empty_cache()
        else:
            pipe.load_models_to_device(self.onload_model_names)
            prompt_embeds = self.encode_prompt(pipe, prompt, pipe.device)

        return {"prompt_embeds": prompt_embeds}


class ZImageUnit_NoiseInitializer(PipelineUnit):
    def __init__(self):
        super().__init__(
            input_params=("height", "width", "seed", "rand_device"),
            output_params=("noise",),
        )

    def process(self, pipe: ZImagePipeline, height, width, seed, rand_device):
        # noise = pipe.generate_noise((1, 16, height//8, width//8), seed=seed, rand_device=rand_device, rand_torch_dtype=pipe.torch_dtype)
        noise = pipe.generate_noise(
            (1, 3, height, width), 
            seed=seed, 
            rand_device=rand_device, 
            rand_torch_dtype=torch.float32,
            torch_dtype=torch.float32
        )
        return {"noise": noise}


class ZImageUnit_InputImageEmbedder(PipelineUnit):
    def __init__(self):
        super().__init__(
            input_params=("input_image", "noise"),
            output_params=("latents", "input_latents"),
            onload_model_names=()
            # onload_model_names=("vae_encoder",)
        )

    def process(self, pipe: ZImagePipeline, input_image, noise):
        if input_image is None:
            return {"latents": noise, "input_latents": None}
        # pipe.load_models_to_device(['vae'])
        image = pipe.preprocess_image(input_image)
        # input_latents = pipe.vae_encoder(image)
        input_latents = image.to(dtype=noise.dtype, device=noise.device)

        if pipe.scheduler.training:
            return {"latents": noise, "input_latents": input_latents}
        else:
            latents = pipe.scheduler.add_noise(input_latents, noise, timestep=pipe.scheduler.timesteps[0])
            return {"latents": latents, "input_latents": input_latents}


def model_fn_z_image(
    dit: ZImageDiT,
    latents=None,
    timestep=None,
    prompt_embeds=None,
    use_gradient_checkpointing=False,
    use_gradient_checkpointing_offload=False,
    **kwargs,
):
    latents = [rearrange(latents.to(next(dit.parameters()).dtype), "B C H W -> C B H W")]
    timestep = (1000 - timestep) / 1000
    model_output = dit(
        latents,
        timestep,
        prompt_embeds,
        patch_size=16,
        use_gradient_checkpointing=use_gradient_checkpointing,
        use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
    )[0][0]
    model_output = -model_output
    model_output = rearrange(model_output, "C B H W -> B C H W")
    return model_output