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()