omini_control/flux_conceptrol_pipeline.py

import torch
from diffusers.pipelines import FluxPipeline
from typing import List, Union, Optional, Dict, Any, Callable
from .transformer import tranformer_forward
from .condition import Condition
from .conceptrol import Conceptrol

from diffusers.pipelines.flux.pipeline_flux import (
    FluxPipelineOutput,
    calculate_shift,
    retrieve_timesteps,
    np,
)

denoising_images = []


def prepare_params(
    prompt: Union[str, List[str]] = None,
    prompt_2: Optional[Union[str, List[str]]] = None,
    height: Optional[int] = 512,
    width: Optional[int] = 512,
    num_inference_steps: int = 28,
    timesteps: List[int] = None,
    guidance_scale: float = 3.5,
    num_images_per_prompt: Optional[int] = 1,
    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
    callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
    callback_on_step_end_tensor_inputs: List[str] = ["latents"],
    max_sequence_length: int = 512,
    **kwargs: dict,
):
    return (
        prompt,
        prompt_2,
        height,
        width,
        num_inference_steps,
        timesteps,
        guidance_scale,
        num_images_per_prompt,
        generator,
        latents,
        prompt_embeds,
        pooled_prompt_embeds,
        output_type,
        return_dict,
        joint_attention_kwargs,
        callback_on_step_end,
        callback_on_step_end_tensor_inputs,
        max_sequence_length,
    )


def seed_everything(seed: int = 42):
    torch.backends.cudnn.deterministic = True
    torch.manual_seed(seed)
    np.random.seed(seed)


def set_scale(pipe, condition_scale):
    for name, module in pipe.transformer.named_modules():
        if not name.endswith(".attn"):
            continue
        module.c_factor = torch.ones(1, 1) * condition_scale


class FluxConceptrolPipeline(FluxPipeline):

    def find_subsequence(self, text, sub):
        sub_len = len(sub)
        for i in range(len(text) - sub_len + 1):
            if text[i : i + sub_len] == sub:
                return i, i + sub_len  # Return start and end indices
        return None

    def locate_subject(self, prompt, subject, max_length=512):
        text_inputs = self.tokenizer_2.tokenize(
            prompt,
            padding="max_length",
            max_length=max_length,
            truncation=True,
            return_length=False,
            return_overflowing_tokens=False,
            return_tensors="pt",
        )
        subject_inputs = self.tokenizer_2.tokenize(
            subject,
            truncation=True,
            return_length=False,
            return_overflowing_tokens=False,
            return_tensors="pt",
        )
        print(f"Locate {subject}", self.find_subsequence(text_inputs, subject_inputs))
        return self.find_subsequence(text_inputs, subject_inputs)

        text_input_ids = text_inputs
        return (
            text_input_ids.index(subject_inputs[0]),
            text_input_ids.index(subject_inputs[-1]) + 1,
        )

    def load_conceptrol(self, conceptrol):
        self.conceptrol = conceptrol

    @torch.no_grad()
    def __call__(
        self,
        image=None,
        model_config: Optional[Dict[str, Any]] = {},
        condition_scale: float = 1.0,
        subject: Optional[str] = None,
        control_guidance_start: float = 0.0,
        control_guidance_end: float = 1.0,
        conceptrol: Conceptrol = None,
        seed: int = 42,
        **params: dict,
    ):
        seed_everything(seed)

        if conceptrol is None:
            if not hasattr(self, "conceptrol"):
                raise ValueError("Default conceptrol not loaded. Please call load_conceptrol() first.")
            conceptrol = self.conceptrol

        conditions = [Condition("subject", image.convert("RGB").resize((512, 512)))]
        if condition_scale != 1:
            for name, module in self.transformer.named_modules():
                if not name.endswith(".attn"):
                    continue
                module.c_factor = torch.ones(1, 1) * condition_scale

        (
            prompt,
            prompt_2,
            height,
            width,
            num_inference_steps,
            timesteps,
            guidance_scale,
            num_images_per_prompt,
            generator,
            latents,
            prompt_embeds,
            pooled_prompt_embeds,
            output_type,
            return_dict,
            joint_attention_kwargs,
            callback_on_step_end,
            callback_on_step_end_tensor_inputs,
            max_sequence_length,
        ) = prepare_params(**params)

        if subject is not None:
            textual_concept_idx = self.locate_subject(params["prompt"], subject)
        else:
            raise ValueError("Subject has to be provided")

        if textual_concept_idx is None:
            raise ValueError("Textual concept idx has to be provided")

        conceptrol.register(textual_concept_idx)

        height = height or self.default_sample_size * self.vae_scale_factor
        width = width or self.default_sample_size * self.vae_scale_factor

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt,
            prompt_2,
            height,
            width,
            prompt_embeds=prompt_embeds,
            pooled_prompt_embeds=pooled_prompt_embeds,
            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
            max_sequence_length=max_sequence_length,
        )

        self._guidance_scale = guidance_scale
        self._joint_attention_kwargs = joint_attention_kwargs
        self._interrupt = False

        # 2. Define call parameters
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        device = self._execution_device

        lora_scale = (
            self.joint_attention_kwargs.get("scale", None)
            if self.joint_attention_kwargs is not None
            else None
        )
        (
            prompt_embeds,
            pooled_prompt_embeds,
            text_ids,
        ) = self.encode_prompt(
            prompt=prompt,
            prompt_2=prompt_2,
            prompt_embeds=prompt_embeds,
            pooled_prompt_embeds=pooled_prompt_embeds,
            device=device,
            num_images_per_prompt=num_images_per_prompt,
            max_sequence_length=max_sequence_length,
            lora_scale=lora_scale,
        )

        # 4. Prepare latent variables
        num_channels_latents = self.transformer.config.in_channels // 4
        latents, latent_image_ids = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
        )

        # 4.1. Prepare conditions
        condition_latents, condition_ids, condition_type_ids = ([] for _ in range(3))
        use_condition = conditions is not None or []
        if use_condition:
            assert len(conditions) <= 1, "Only one condition is supported for now."
            self.set_adapters(conditions[0].condition_type)
            for condition in conditions:
                tokens, ids, type_id = condition.encode(self)
                condition_latents.append(tokens)  # [batch_size, token_n, token_dim]
                condition_ids.append(ids)  # [token_n, id_dim(3)]
                condition_type_ids.append(type_id)  # [token_n, 1]
            condition_latents = torch.cat(condition_latents, dim=1)
            condition_ids = torch.cat(condition_ids, dim=0)
            if condition.condition_type == "subject":
                condition_ids[:, 2] += width // 16
            condition_type_ids = torch.cat(condition_type_ids, dim=0)

        # 5. Prepare timesteps
        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
        image_seq_len = latents.shape[1]
        mu = calculate_shift(
            image_seq_len,
            self.scheduler.config.base_image_seq_len,
            self.scheduler.config.max_image_seq_len,
            self.scheduler.config.base_shift,
            self.scheduler.config.max_shift,
        )
        timesteps, num_inference_steps = retrieve_timesteps(
            self.scheduler,
            num_inference_steps,
            device,
            timesteps,
            sigmas,
            mu=mu,
        )
        num_warmup_steps = max(
            len(timesteps) - num_inference_steps * self.scheduler.order, 0
        )
        self._num_timesteps = len(timesteps)

        # 6. Denoising loop
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                if (i / len(timesteps) < control_guidance_start) or (
                    (i + 1) / len(timesteps) > control_guidance_end
                ):
                    set_scale(self, 0.5)  # Warmup required for the first few steps
                else:
                    set_scale(self, condition_scale)
                if self.interrupt:
                    continue

                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
                timestep = t.expand(latents.shape[0]).to(latents.dtype)

                # handle guidance
                if self.transformer.config.guidance_embeds:
                    guidance = torch.tensor([guidance_scale], device=device)
                    guidance = guidance.expand(latents.shape[0])
                else:
                    guidance = None
                noise_pred = tranformer_forward(
                    self.transformer,
                    model_config=model_config,
                    conceptrol=conceptrol,
                    # Inputs of the condition (new feature)
                    condition_latents=condition_latents if use_condition else None,
                    condition_ids=condition_ids if use_condition else None,
                    condition_type_ids=condition_type_ids if use_condition else None,
                    # Inputs to the original transformer
                    hidden_states=latents,
                    timestep=timestep / 1000,
                    guidance=guidance,
                    pooled_projections=pooled_prompt_embeds,
                    encoder_hidden_states=prompt_embeds,
                    txt_ids=text_ids,
                    img_ids=latent_image_ids,
                    joint_attention_kwargs=self.joint_attention_kwargs,
                    return_dict=False,
                )[0]

                # compute the previous noisy sample x_t -> x_t-1
                latents_dtype = latents.dtype
                latents = self.scheduler.step(
                    noise_pred, t, latents, return_dict=False
                )[0]

                if latents.dtype != latents_dtype:
                    if torch.backends.mps.is_available():
                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
                        latents = latents.to(latents_dtype)

                if callback_on_step_end is not None:
                    callback_kwargs = {}
                    for k in callback_on_step_end_tensor_inputs:
                        callback_kwargs[k] = locals()[k]
                    callback_outputs = callback_on_step_end(
                        self, latents, callback_kwargs
                    )

                    global denoising_images
                    denoising_images.append(callback_outputs)

                # call the callback, if provided
                if i == len(timesteps) - 1 or (
                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
                ):
                    progress_bar.update()

        if output_type == "latent":
            image = latents

        else:
            latents = self._unpack_latents(
                latents, height, width, self.vae_scale_factor
            )
            latents = (
                latents / self.vae.config.scaling_factor
            ) + self.vae.config.shift_factor
            image = self.vae.decode(latents, return_dict=False)[0]
            image = self.image_processor.postprocess(image, output_type=output_type)

        # Offload all models
        self.maybe_free_model_hooks()

        if condition_scale != 1:
            for name, module in self.transformer.named_modules():
                if not name.endswith(".attn"):
                    continue
                del module.c_factor

        if not return_dict:
            return (image,)

        return FluxPipelineOutput(images=image)