Update PixelFlow-T2I/pipeline.py

PixelFlow-T2I/pipeline.py

@@ -12,14 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+import inspect
+
 import importlib
 import json
 import math
 import sys
 from pathlib import Path
-from typing import Any, List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union, Any
 
-import numpy as np
 import torch
 import torch.nn.functional as F
 from einops import rearrange
@@ -27,8 +30,6 @@ from einops import rearrange
 from diffusers.image_processor import VaeImageProcessor
 from diffusers.models.embeddings import get_2d_rotary_pos_embed
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
-from diffusers.schedulers import KarrasDiffusionSchedulers
-from diffusers.utils import replace_example_docstring
 from diffusers.utils.torch_utils import randn_tensor
 
 
@@ -38,29 +39,19 @@ EXAMPLE_DOC_STRING = """
     Examples:
         ```py
         >>> from pathlib import Path
+        >>> import sys
         >>> import torch
-        >>> from diffusers import DiffusionPipeline
 
         >>> model_dir = Path("./PixelFlow-T2I").resolve()
-        >>> pipe = DiffusionPipeline.from_pretrained(
+        >>> sys.path.insert(0, str(model_dir))
+        >>> from pipeline import PixelFlowT2IPipeline
+
+        >>> pipe = PixelFlowT2IPipeline.from_pretrained(
         ...     str(model_dir),
         ...     local_files_only=True,
-        ...     custom_pipeline=str(model_dir / "pipeline.py"),
-        ...     trust_remote_code=True,
         ...     torch_dtype=torch.bfloat16,
         ... )
-        >>> pipe = pipe.to("cuda")
-
-        >>> generator = torch.Generator(device="cuda").manual_seed(42)
-        >>> image = pipe(
-        ...     prompt="A golden retriever playing in a sunny garden",
-        ...     height=1024,
-        ...     width=1024,
-        ...     num_inference_steps=[10, 10, 10, 10],
-        ...     guidance_scale=4.0,
-        ...     generator=generator,
-        ... ).images[0]
-        >>> image.save("demo.png")
+        >>> pipe.to("cuda")
         ```
 """
 
@@ -73,14 +64,29 @@ class PixelFlowT2IPipeline(DiffusionPipeline):
     Parameters:
         transformer ([`PixelFlowTransformer2DModel`]):
             Text-conditioned PixelFlow transformer operating in pixel space.
-        scheduler ([`PixelFlowScheduler`] or [`KarrasDiffusionSchedulers`]):
-            Multi-stage flow scheduler used by PixelFlow cascade denoising.
+        scheduler ([`PixelFlowScheduler`]):
+            Multi-stage flow scheduler used by PixelFlow.
         text_encoder ([`T5EncoderModel`], *optional*):
             Text encoder used to embed prompts.
         tokenizer ([`T5Tokenizer`], *optional*):
             Tokenizer paired with the text encoder.
     """
 
+    @staticmethod
+    def prepare_extra_step_kwargs(
+        scheduler,
+        generator=None,
+        eta: float | None = None,
+    ):
+        kwargs = {}
+        step_params = set(inspect.signature(scheduler.step).parameters.keys())
+        if "generator" in step_params:
+            kwargs["generator"] = generator
+        if eta is not None and "eta" in step_params:
+            kwargs["eta"] = eta
+        return kwargs
+
+
     model_cpu_offload_seq = "text_encoder->transformer"
     _optional_components = ["text_encoder", "tokenizer"]
 
@@ -101,7 +107,29 @@ class PixelFlowT2IPipeline(DiffusionPipeline):
         )
         self.image_processor = VaeImageProcessor(vae_scale_factor=1, do_normalize=False)
         self.max_token_length = max_token_length
-        self.set_progress_bar_config(disable=False)
+
+    @staticmethod
+    def _prepare_generator(
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]],
+        device: torch.device,
+    ) -> Optional[Union[torch.Generator, List[torch.Generator]]]:
+        if generator is None:
+            return None
+        if isinstance(generator, list):
+            return [PixelFlowT2IPipeline._prepare_generator(item, device) for item in generator]
+
+        gen_device = getattr(generator, "device", torch.device("cpu"))
+        if gen_device.type == device.type:
+            return generator
+        new_generator = torch.Generator(device=device)
+        seed = int(generator.initial_seed())
+        return new_generator.manual_seed(seed)
+
+    def _latent_device(self) -> torch.device:
+        transformer_device = getattr(self.transformer, "device", None)
+        if transformer_device is not None:
+            return transformer_device
+        return self._execution_device
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path=None, subfolder=None, **kwargs):
@@ -136,7 +164,6 @@ class PixelFlowT2IPipeline(DiffusionPipeline):
                 variant = variant / subfolder
 
         model_kwargs = dict(kwargs)
-        model_kwargs.pop("trust_remote_code", None)
         scheduler_kwargs = model_kwargs.pop("scheduler_kwargs", {})
         inserted = []
 
@@ -154,14 +181,15 @@ class PixelFlowT2IPipeline(DiffusionPipeline):
             transformer_cls = getattr(importlib.import_module("transformer_pixelflow"), "PixelFlowTransformer2DModel")
             transformer = transformer_cls.from_pretrained(str(transformer_dir), **model_kwargs)
 
-            scheduler_dir = variant / "scheduler"
-            if not (scheduler_dir / "scheduler_config.json").exists():
-                raise FileNotFoundError(f"Expected scheduler config in {scheduler_dir}")
+            scheduling_py = variant / "scheduling_pixelflow.py"
+            scheduler_cfg_dir = variant / "scheduler"
+            if not scheduling_py.is_file() or not (scheduler_cfg_dir / "scheduler_config.json").exists():
+                raise FileNotFoundError(f"Expected scheduler module at {scheduling_py} and config in {scheduler_cfg_dir}")
 
-            _ensure_path(str(scheduler_dir))
+            _ensure_path(str(variant.resolve()))
             scheduler_cls = getattr(importlib.import_module("scheduling_pixelflow"), "PixelFlowScheduler")
             try:
-                scheduler = scheduler_cls.from_pretrained(str(scheduler_dir), **scheduler_kwargs)
+                scheduler = scheduler_cls.from_pretrained(str(scheduler_cfg_dir), **scheduler_kwargs)
             except Exception:
                 scheduler = scheduler_cls(**scheduler_kwargs)
 
@@ -187,6 +215,70 @@ class PixelFlowT2IPipeline(DiffusionPipeline):
                 if comp_path in sys.path:
                     sys.path.remove(comp_path)
 
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        model_kwargs = dict(kwargs)
+        transformer_subfolder = model_kwargs.pop("transformer_subfolder", None)
+        scheduler_subfolder = model_kwargs.pop("scheduler_subfolder", None)
+        text_encoder_subfolder = model_kwargs.pop("text_encoder_subfolder", None)
+        tokenizer_subfolder = model_kwargs.pop("tokenizer_subfolder", None)
+        scheduler_kwargs = model_kwargs.pop("scheduler_kwargs", {})
+        base_path = Path(pretrained_model_name_or_path)
+
+        if transformer_subfolder is None and (base_path / "transformer").exists():
+            transformer_subfolder = "transformer"
+        if scheduler_subfolder is None and (base_path / "scheduler").exists():
+            scheduler_subfolder = "scheduler"
+        if text_encoder_subfolder is None and (base_path / "text_encoder").exists():
+            text_encoder_subfolder = "text_encoder"
+        if tokenizer_subfolder is None and (base_path / "tokenizer").exists():
+            tokenizer_subfolder = "tokenizer"
+
+        try:
+            return super().from_pretrained(pretrained_model_name_or_path, **kwargs)
+        except Exception:
+            if transformer_subfolder is not None:
+                transformer_path = str(base_path / transformer_subfolder)
+            else:
+                transformer_path = pretrained_model_name_or_path
+
+            transformer = PixelFlowTransformer2DModel.from_pretrained(transformer_path, **model_kwargs)
+            try:
+                scheduler = PixelFlowScheduler.from_pretrained(
+                    pretrained_model_name_or_path,
+                    subfolder=scheduler_subfolder,
+                    **scheduler_kwargs,
+                )
+            except Exception:
+                scheduler = PixelFlowScheduler(**scheduler_kwargs)
+
+            text_encoder = None
+            tokenizer = None
+            if text_encoder_subfolder is not None and (base_path / text_encoder_subfolder / "config.json").exists():
+                from transformers import T5EncoderModel, T5Tokenizer
+
+                text_encoder = T5EncoderModel.from_pretrained(
+                    str(base_path / text_encoder_subfolder),
+                    **model_kwargs,
+                )
+                tokenizer = T5Tokenizer.from_pretrained(str(base_path / tokenizer_subfolder))
+
+            if text_encoder is None and tokenizer is None:
+                text_encoder_name = cls._read_text_encoder_name(base_path)
+                from transformers import T5EncoderModel, T5Tokenizer
+
+                text_encoder = T5EncoderModel.from_pretrained(text_encoder_name, **model_kwargs)
+                tokenizer = T5Tokenizer.from_pretrained(text_encoder_name)
+
+            pipe = cls(
+                transformer=transformer,
+                scheduler=scheduler,
+                text_encoder=text_encoder,
+                tokenizer=tokenizer,
+            )
+            if hasattr(pipe, "register_to_config"):
+                pipe.register_to_config(_name_or_path=str(base_path))
+            return pipe
+
     @staticmethod
     def _read_text_encoder_name(variant_path: Path) -> str:
         metadata_path = variant_path / "conversion_metadata.json"
@@ -243,11 +335,12 @@ class PixelFlowT2IPipeline(DiffusionPipeline):
         init_factor = 2 ** (self.scheduler.num_stages - 1)
         coarse_height = height // init_factor
         coarse_width = width // init_factor
+        latent_dtype = getattr(self.transformer, "dtype", torch.float32)
         latents = randn_tensor(
             (batch_size, 3, coarse_height, coarse_width),
             generator=generator,
             device=device,
-            dtype=torch.float32,
+            dtype=latent_dtype,
         )
         return latents, coarse_height, coarse_width
 
@@ -257,15 +350,23 @@ class PixelFlowT2IPipeline(DiffusionPipeline):
         channels: int,
         height: int,
         width: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
         eps: float = 1e-6,
     ) -> torch.Tensor:
         gamma = self.scheduler.gamma
-        dist = torch.distributions.multivariate_normal.MultivariateNormal(
-            torch.zeros(4),
-            torch.eye(4) * (1 - gamma) + torch.ones(4, 4) * gamma + eps * torch.eye(4),
-        )
+        cov = torch.eye(4, dtype=torch.float32) * (1 - gamma) + torch.ones(4, 4, dtype=torch.float32) * gamma
+        cov = cov + eps * torch.eye(4, dtype=torch.float32)
+        chol = torch.linalg.cholesky(cov).to(device=device, dtype=dtype)
         block_number = batch_size * channels * (height // 2) * (width // 2)
-        noise = torch.stack([dist.sample() for _ in range(block_number)])
+        standard = randn_tensor(
+            (block_number, 4),
+            generator=generator,
+            device=device,
+            dtype=dtype,
+        )
+        noise = standard @ chol.T
         return rearrange(
             noise,
             "(b c h w) (p q) -> b c (h p) (w q)",
@@ -284,6 +385,7 @@ class PixelFlowT2IPipeline(DiffusionPipeline):
         height: int,
         width: int,
         device: torch.device,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
     ) -> torch.Tensor:
         latents = F.interpolate(latents, size=(height, width), mode="nearest")
         original_start_t = self.scheduler.original_start_t[stage_idx]
@@ -291,8 +393,12 @@ class PixelFlowT2IPipeline(DiffusionPipeline):
         alpha = 1 / (math.sqrt(1 - (1 / gamma)) * (1 - original_start_t) + original_start_t)
         beta = alpha * (1 - original_start_t) / math.sqrt(-gamma)
 
-        noise = self._sample_block_noise(*latents.shape)
-        noise = noise.to(device=device, dtype=latents.dtype)
+        noise = self._sample_block_noise(
+            *latents.shape,
+            device=device,
+            dtype=latents.dtype,
+            generator=generator,
+        )
         return alpha * latents + beta * noise
 
     def _prepare_rope_pos_embed(self, latents: torch.Tensor, device: torch.device) -> torch.Tensor:
@@ -415,7 +521,6 @@ class PixelFlowT2IPipeline(DiffusionPipeline):
         return prompt_embeds, prompt_attention_mask
 
     @torch.inference_mode()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
     def __call__(
         self,
         prompt: Union[str, List[str]],
@@ -433,9 +538,6 @@ class PixelFlowT2IPipeline(DiffusionPipeline):
         r"""
         Generate text-to-image samples with PixelFlow.
 
-        Examples:
-            <!-- this section is replaced by replace_example_docstring -->
-
         Args:
             prompt (`str` or `list[str]`):
                 Text prompt(s) describing the desired image.
@@ -470,65 +572,71 @@ class PixelFlowT2IPipeline(DiffusionPipeline):
         width = int(width or default_size)
         self.check_inputs(prompt_list, height, width, num_inference_steps, output_type, negative_prompt)
 
-        device = self.transformer.device
-        text_encoder_device = self.text_encoder.device if self.text_encoder is not None else device
+        device = self._execution_device
+        latent_device = self._latent_device()
         do_classifier_free_guidance = guidance_scale > 1.0
         stage_steps = self._normalize_stage_steps(num_inference_steps)
         batch_size = len(prompt_list)
+        generator = self._prepare_generator(generator, latent_device)
 
         prompt_embeds, prompt_attention_mask = self.encode_prompt(
             prompt_list,
-            text_encoder_device,
+            device,
             num_images_per_prompt=num_images_per_prompt,
             do_classifier_free_guidance=do_classifier_free_guidance,
             negative_prompt=negative_prompt,
         )
-        prompt_embeds = prompt_embeds.to(device)
-        prompt_attention_mask = prompt_attention_mask.to(device)
 
         latents, height, width = self.prepare_latents(
             batch_size * num_images_per_prompt,
             height,
             width,
-            device,
+            latent_device,
             generator,
         )
-        size_tensor = torch.tensor([latents.shape[-1] // self.transformer.patch_size], dtype=torch.int32, device=device)
+        latents = latents.to(device=latent_device, dtype=self.transformer.dtype)
+        prompt_embeds = prompt_embeds.to(device=latent_device)
+        prompt_attention_mask = prompt_attention_mask.to(device=latent_device)
+        size_tensor = torch.tensor([latents.shape[-1] // self.transformer.patch_size], dtype=torch.int32, device=latent_device)
 
-        autocast_enabled = device.type == "cuda"
+        autocast_enabled = latent_device.type == "cuda"
         autocast_dtype = torch.bfloat16 if autocast_enabled else torch.float32
 
-        with self.progress_bar(total=sum(stage_steps)) as progress_bar:
-            for stage_idx in range(self.scheduler.num_stages):
-                self.scheduler.set_timesteps(stage_steps[stage_idx], stage_idx, device=device, shift=shift)
-                timesteps = self.scheduler.Timesteps
-
-                if stage_idx > 0:
-                    height, width = height * 2, width * 2
-                    latents = self._upsample_latents_for_stage(latents, stage_idx, height, width, device)
-                    size_tensor = torch.tensor([latents.shape[-1] // self.transformer.patch_size], dtype=torch.int32, device=device)
-
-                rope_pos = self._prepare_rope_pos_embed(latents, device)
-
-                for timestep in timesteps:
-                    latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-                    timestep_batch = timestep.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
-                    with torch.autocast(device.type, enabled=autocast_enabled, dtype=autocast_dtype):
-                        noise_pred = self.transformer(
-                            latent_model_input,
-                            encoder_hidden_states=prompt_embeds,
-                            encoder_attention_mask=prompt_attention_mask,
-                            timestep=timestep_batch,
-                            latent_size=size_tensor,
-                            pos_embed=rope_pos,
-                        ).sample
-
-                    if do_classifier_free_guidance:
-                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-                    latents = self.scheduler.step(model_output=noise_pred, sample=latents).prev_sample
-                    progress_bar.update()
+        extra_step_kwargs = self.prepare_extra_step_kwargs(self.scheduler, generator=generator)
+
+        for stage_idx in range(self.scheduler.num_stages):
+            self.scheduler.set_timesteps(stage_steps[stage_idx], stage_idx, device=latent_device, shift=shift)
+            timesteps = self.scheduler.Timesteps
+
+            if stage_idx > 0:
+                height, width = height * 2, width * 2
+                latents = self._upsample_latents_for_stage(
+                    latents, stage_idx, height, width, latent_device, generator=generator
+                )
+                latents = latents.to(dtype=self.transformer.dtype)
+                size_tensor = torch.tensor([latents.shape[-1] // self.transformer.patch_size], dtype=torch.int32, device=latent_device)
+
+            rope_pos = self._prepare_rope_pos_embed(latents, latent_device)
+
+            for timestep in timesteps:
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = latent_model_input.to(device=latent_device, dtype=self.transformer.dtype)
+                timestep_batch = timestep.expand(latent_model_input.shape[0]).to(device=latent_device, dtype=self.transformer.dtype)
+                with torch.autocast(latent_device.type, enabled=autocast_enabled, dtype=autocast_dtype):
+                    noise_pred = self.transformer(
+                        latent_model_input,
+                        encoder_hidden_states=prompt_embeds,
+                        encoder_attention_mask=prompt_attention_mask,
+                        timestep=timestep_batch,
+                        latent_size=size_tensor,
+                        pos_embed=rope_pos,
+                    ).sample
+
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                latents = self.scheduler.step(model_output=noise_pred, sample=latents, **extra_step_kwargs).prev_sample
 
         image = self.decode_latents(latents, output_type=output_type)
         self.maybe_free_model_hooks()