lrm/flux/trainer/models/flux_preference_model.py

import os
import sys
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path

import torch
from PIL import Image
from torch import nn
from torchvision import transforms
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast

# Ensure sibling project modules (information_related_to_flux) are importable
_PROJECT_ROOT = Path(__file__).resolve().parents[3]
if str(_PROJECT_ROOT) not in sys.path:
    sys.path.append(str(_PROJECT_ROOT))

from information_related_to_flux.dit import FluxTransformer2DModel
from information_related_to_flux.pipeline import FluxPipeline
from trainer.models.base_model import BaseModelConfig


@dataclass
class FluxPreferenceModelConfig(BaseModelConfig):
    _target_: str = "trainer.models.flux_preference_model.FluxPreferenceModel"
    pretrained_model_name_or_path: str = "black-forest-labs/FLUX.1-schnell"
    pretrained_vae_name_or_path: str = "black-forest-labs/FLUX.1-schnell"
    projection_dim: int = 1024
    text_embed_dim: int = 768
    logit_scale_init_value: float = 2.6592
    freeze_text_encoder: bool = False
    guidance_scale: float = 0.0
    noise_offset: bool = False
    noise_offset_coeff: float = 0.05
    max_sequence_length: int = 512
    image_size: int = 1024


class FluxPreferenceModel(nn.Module):
    def __init__(self, cfg: FluxPreferenceModelConfig):
        super().__init__()
        self.cfg = cfg

        offline_mode = os.getenv("HF_HUB_OFFLINE", "0").strip().lower() in {"1", "true", "yes", "on"}
        cache_dir = os.getenv("HF_HUB_CACHE") or os.getenv("HUGGINGFACE_HUB_CACHE")
        pretrained_kwargs = {
            "local_files_only": offline_mode,
        }
        if cache_dir:
            pretrained_kwargs["cache_dir"] = cache_dir

        # Keep weights in the requested mixed precision to avoid fp32 VRAM blowups.
        precision = os.getenv("ACCELERATE_MIXED_PRECISION", "").strip().lower()
        model_dtype = None
        if precision == "bf16":
            model_dtype = torch.bfloat16
        elif precision == "fp16":
            model_dtype = torch.float16

        module_load_kwargs = dict(pretrained_kwargs)
        if model_dtype is not None:
            module_load_kwargs["torch_dtype"] = model_dtype

        self.vae = AutoencoderKL.from_pretrained(
            cfg.pretrained_vae_name_or_path,
            subfolder="vae",
            **module_load_kwargs,
        )
        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
            cfg.pretrained_model_name_or_path,
            subfolder="scheduler",
            **pretrained_kwargs,
        )
        self.transformer = FluxTransformer2DModel.from_pretrained(
            cfg.pretrained_model_name_or_path,
            subfolder="transformer",
            **module_load_kwargs,
        )
        self.tokenizer = CLIPTokenizer.from_pretrained(
            cfg.pretrained_model_name_or_path,
            subfolder="tokenizer",
            **pretrained_kwargs,
        )
        self.tokenizer_2 = T5TokenizerFast.from_pretrained(
            cfg.pretrained_model_name_or_path,
            subfolder="tokenizer_2",
            **pretrained_kwargs,
        )
        self.text_encoder = CLIPTextModel.from_pretrained(
            cfg.pretrained_model_name_or_path,
            subfolder="text_encoder",
            **module_load_kwargs,
        )
        self.text_encoder_2 = T5EncoderModel.from_pretrained(
            cfg.pretrained_model_name_or_path,
            subfolder="text_encoder_2",
            **module_load_kwargs,
        )

        self.vae.requires_grad_(False)
        if cfg.freeze_text_encoder:
            self.text_encoder.requires_grad_(False)
            self.text_encoder_2.requires_grad_(False)

        text_in_dim = self.text_encoder.config.hidden_size
        image_in_dim = self.transformer.config.in_channels

        self.text_projection = nn.Linear(text_in_dim, cfg.projection_dim, bias=False)
        self.visual_projection = nn.Linear(image_in_dim, cfg.projection_dim, bias=False)
        nn.init.normal_(self.text_projection.weight, std=0.02)
        nn.init.normal_(self.visual_projection.weight, std=0.02)

        self.logit_scale = nn.Parameter(torch.ones([]) * cfg.logit_scale_init_value)

        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
        self.height = cfg.image_size
        self.width = cfg.image_size
        self.val_transform = transforms.Compose(
            [
                transforms.Resize((self.height, self.width), interpolation=transforms.InterpolationMode.BILINEAR),
                transforms.ToTensor(),
                transforms.Normalize([0.5], [0.5]),
            ]
        )

    def _get_sigmas_from_indices(self, timestep_indices: torch.Tensor, n_dim: int, dtype: torch.dtype):
        all_sigmas = self.scheduler.sigmas.to(device=timestep_indices.device, dtype=dtype)
        max_index = all_sigmas.shape[0] - 1
        timestep_indices = timestep_indices.clamp(0, max_index).long()
        sigma = all_sigmas[timestep_indices].flatten()
        while len(sigma.shape) < n_dim:
            sigma = sigma.unsqueeze(-1)
        return sigma

    def _encode_prompt(self, text_input_ids: torch.Tensor, text_input_ids_2: torch.Tensor):
        clip_out = self.text_encoder(text_input_ids, output_hidden_states=False)
        pooled_prompt_embeds = clip_out.pooler_output
        prompt_embeds = self.text_encoder_2(text_input_ids_2, output_hidden_states=False)[0]

        pooled_prompt_embeds = pooled_prompt_embeds.to(dtype=self.text_encoder.dtype, device=text_input_ids.device)
        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=text_input_ids_2.device)

        text_ids = torch.zeros(prompt_embeds.shape[1], 3, device=prompt_embeds.device, dtype=prompt_embeds.dtype)
        text_features = self.text_projection(pooled_prompt_embeds)
        return prompt_embeds, pooled_prompt_embeds, text_ids, text_features

    def _encode_images(self, image_inputs: torch.Tensor):
        vae_param = next(self.vae.parameters())
        image_inputs = image_inputs.to(device=vae_param.device, dtype=vae_param.dtype)
        with torch.no_grad():
            latents = self.vae.encode(image_inputs).latent_dist.sample()
            latents = (latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
        return latents

    def get_image_features(
        self,
        encoder_hidden_states: torch.Tensor,
        pooled_prompt_embeds: torch.Tensor,
        text_ids: torch.Tensor,
        image_inputs: torch.Tensor,
        time_cond: torch.Tensor,
        generator=None,
    ):
        latents = self._encode_images(image_inputs)

        if generator is not None:
            noise = torch.randn(latents.size(), generator=generator, dtype=latents.dtype, device=latents.device)
        else:
            noise = torch.randn_like(latents)

        if self.cfg.noise_offset:
            noise = noise + self.cfg.noise_offset_coeff * torch.randn(
                (latents.shape[0], latents.shape[1], 1, 1),
                device=latents.device,
                dtype=latents.dtype,
            )

        sigmas = self._get_sigmas_from_indices(time_cond, n_dim=latents.ndim, dtype=latents.dtype)
        noisy_latents = (1.0 - sigmas) * latents + sigmas * noise

        packed_noisy_latents = FluxPipeline._pack_latents(
            noisy_latents,
            batch_size=latents.shape[0],
            num_channels_latents=latents.shape[1],
            height=latents.shape[2],
            width=latents.shape[3],
        )

        latent_image_ids = FluxPipeline._prepare_latent_image_ids(
            latents.shape[0],
            latents.shape[2] // 2,
            latents.shape[3] // 2,
            latents.device,
            latents.dtype,
        )

        guidance = None
        if self.transformer.config.guidance_embeds:
            guidance = torch.full(
                (latents.shape[0],),
                self.cfg.guidance_scale,
                device=latents.device,
                dtype=latents.dtype,
            )

        scheduler_timesteps = self.scheduler.timesteps.to(device=time_cond.device)
        timestep = scheduler_timesteps[time_cond.long()].to(device=latents.device, dtype=latents.dtype)
        model_pred = self.transformer(
            hidden_states=packed_noisy_latents,
            timestep=timestep / 1000,
            guidance=guidance,
            pooled_projections=pooled_prompt_embeds,
            encoder_hidden_states=encoder_hidden_states,
            txt_ids=text_ids,
            img_ids=latent_image_ids,
            return_dict=False,
        )[0]

        pooled_tokens = model_pred.mean(dim=1)
        image_features = self.visual_projection(pooled_tokens)
        return image_features

    def forward(self, text_input_ids, text_input_ids_2, image_inputs, time_cond, generator=None):
        n_prompts = text_input_ids.shape[0]
        n_images = image_inputs.shape[0]

        encoder_hidden_states, pooled_prompt_embeds, text_ids, text_features = self._encode_prompt(
            text_input_ids,
            text_input_ids_2,
        )

        if n_images == 2 * n_prompts:
            encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states], dim=0)
            pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, pooled_prompt_embeds], dim=0)

        image_features = self.get_image_features(
            encoder_hidden_states=encoder_hidden_states,
            pooled_prompt_embeds=pooled_prompt_embeds,
            text_ids=text_ids,
            image_inputs=image_inputs,
            time_cond=time_cond,
            generator=generator,
        )

        return text_features, image_features

    def save(self, path):
        self.transformer.save_pretrained(os.path.join(path, "transformer"), safe_serialization=True)
        if not self.cfg.freeze_text_encoder:
            self.text_encoder.save_pretrained(os.path.join(path, "text_encoder"), safe_serialization=True)
            self.text_encoder_2.save_pretrained(os.path.join(path, "text_encoder_2"), safe_serialization=True)

        state_dict = {
            "visual_projection": self.visual_projection.state_dict(),
            "text_projection": self.text_projection.state_dict(),
            "logit_scale": self.logit_scale.data.item(),
        }
        torch.save(state_dict, os.path.join(path, "state_dict.pt"))

    def load(self, path):
        self.transformer = self.transformer.from_pretrained(os.path.join(path, "transformer"))
        if not self.cfg.freeze_text_encoder:
            self.text_encoder = self.text_encoder.from_pretrained(os.path.join(path, "text_encoder"))
            self.text_encoder_2 = self.text_encoder_2.from_pretrained(os.path.join(path, "text_encoder_2"))

        state_dict = torch.load(os.path.join(path, "state_dict.pt"), map_location="cpu")
        self.visual_projection.load_state_dict(state_dict["visual_projection"])
        self.text_projection.load_state_dict(state_dict["text_projection"])
        self.logit_scale.data = torch.tensor(state_dict["logit_scale"])

    def encode_prompt(self, prompt):
        text_input_ids = self.tokenizer(
            prompt,
            padding="max_length",
            max_length=self.tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt",
        ).input_ids
        text_input_ids_2 = self.tokenizer_2(
            prompt,
            padding="max_length",
            max_length=self.cfg.max_sequence_length,
            truncation=True,
            return_tensors="pt",
        ).input_ids
        return text_input_ids, text_input_ids_2

    def preprocess_image(self, images):
        if not isinstance(images, list):
            images = [images]

        image_inputs = []
        for image in images:
            if isinstance(image, dict):
                image = image["bytes"]
            if isinstance(image, bytes):
                image = Image.open(BytesIO(image))
            elif isinstance(image, str):
                image = Image.open(image)
            image = image.convert("RGB")
            image = self.val_transform(image)
            image_inputs.append(image)
        image_inputs = torch.stack(image_inputs, dim=0)
        return image_inputs

    def get_preference_scores(self, prompt, images, timesteps, generator=None):
        image_inputs = self.preprocess_image(images).to(self.vae.device, dtype=self.vae.dtype)
        text_input_ids, text_input_ids_2 = self.encode_prompt(prompt)
        text_input_ids = text_input_ids.to(self.text_encoder.device)
        text_input_ids_2 = text_input_ids_2.to(self.text_encoder_2.device)
        timestep_indices = torch.tensor([timesteps] * image_inputs.shape[0], dtype=torch.long, device=self.vae.device)

        with torch.no_grad():
            text_embs, image_embs = self.forward(
                text_input_ids,
                text_input_ids_2,
                image_inputs,
                timestep_indices,
                generator=generator,
            )

            image_embs = image_embs / torch.norm(image_embs, dim=-1, keepdim=True)
            text_embs = text_embs / torch.norm(text_embs, dim=-1, keepdim=True)

            scores = self.logit_scale.exp() * (text_embs @ image_embs.T)[0]
            probs = torch.softmax(scores, dim=-1)

        return scores.cpu().tolist(), probs.cpu().tolist()