Blame for toolkit/models/diffusion_feature_extraction.py

import math
import torch
import os
from torch import nn
from safetensors.torch import load_file
import torch.nn.functional as F
import torch.utils.checkpoint as ckpt
from diffusers import AutoencoderTiny
from transformers import AutoImageProcessor, AutoModel, SiglipImageProcessor, SiglipVisionModel
import lpips
import weakref
from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
from toolkit.models.base_model import BaseModel
from toolkit.models.sapiens2 import Sapiens2
import huggingface_hub


def _fold_frames_to_batch(x: torch.Tensor) -> torch.Tensor:
    """(B, C, T, H, W) -> (B*T, C, H, W), each sample's frames contiguous -- so the 2D
    feature losses run on EVERY frame of a video instead of only the first."""
    b, c, t, h, w = x.shape
    return x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)


class ResBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
        self.norm1 = nn.GroupNorm(8, out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
        self.norm2 = nn.GroupNorm(8, out_channels)
        self.skip = nn.Conv2d(in_channels, out_channels,
                              1) if in_channels != out_channels else nn.Identity()

    def forward(self, x):
        identity = self.skip(x)
        x = self.conv1(x)
        x = self.norm1(x)
        x = F.silu(x)
        x = self.conv2(x)
        x = self.norm2(x)
        x = F.silu(x + identity)
        return x


class DiffusionFeatureExtractor2(nn.Module):
    def __init__(self, in_channels=32):
        super().__init__()
        self.version = 2

        # Path 1: Upsample to 512x512 (1, 64, 512, 512)
        self.up_path = nn.ModuleList([
            nn.Conv2d(in_channels, 64, 3, padding=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            ResBlock(64, 64),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            ResBlock(64, 64),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            ResBlock(64, 64),
            nn.Conv2d(64, 64, 3, padding=1),
        ])

        # Path 2: Upsample to 256x256 (1, 128, 256, 256)
        self.path2 = nn.ModuleList([
            nn.Conv2d(in_channels, 128, 3, padding=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            ResBlock(128, 128),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            ResBlock(128, 128),
            nn.Conv2d(128, 128, 3, padding=1),
        ])

        # Path 3: Upsample to 128x128 (1, 256, 128, 128)
        self.path3 = nn.ModuleList([
            nn.Conv2d(in_channels, 256, 3, padding=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            ResBlock(256, 256),
            nn.Conv2d(256, 256, 3, padding=1)
        ])

        # Path 4: Original size (1, 512, 64, 64)
        self.path4 = nn.ModuleList([
            nn.Conv2d(in_channels, 512, 3, padding=1),
            ResBlock(512, 512),
            ResBlock(512, 512),
            nn.Conv2d(512, 512, 3, padding=1)
        ])

        # Path 5: Downsample to 32x32 (1, 512, 32, 32)
        self.path5 = nn.ModuleList([
            nn.Conv2d(in_channels, 512, 3, padding=1),
            ResBlock(512, 512),
            nn.AvgPool2d(2),
            ResBlock(512, 512),
            nn.Conv2d(512, 512, 3, padding=1)
        ])

    def forward(self, x):
        outputs = []

        # Path 1: 512x512
        x1 = x
        for layer in self.up_path:
            x1 = layer(x1)
        outputs.append(x1)  # [1, 64, 512, 512]

        # Path 2: 256x256
        x2 = x
        for layer in self.path2:
            x2 = layer(x2)
        outputs.append(x2)  # [1, 128, 256, 256]

        # Path 3: 128x128
        x3 = x
        for layer in self.path3:
            x3 = layer(x3)
        outputs.append(x3)  # [1, 256, 128, 128]

        # Path 4: 64x64
        x4 = x
        for layer in self.path4:
            x4 = layer(x4)
        outputs.append(x4)  # [1, 512, 64, 64]

        # Path 5: 32x32
        x5 = x
        for layer in self.path5:
            x5 = layer(x5)
        outputs.append(x5)  # [1, 512, 32, 32]

        return outputs


class DFEBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)
        self.act = nn.GELU()
        self.proj = nn.Conv2d(channels, channels, 1)

    def forward(self, x):
        x_in = x
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.act(x)
        x = self.proj(x)
        x = x + x_in
        return x


class DiffusionFeatureExtractor(nn.Module):
    def __init__(self, in_channels=16, out_channels=512):
        super().__init__()
        self.version = 1
        num_blocks = 6
        self.conv_in = nn.Conv2d(in_channels, 512, 1)
        self.blocks = nn.ModuleList([DFEBlock(512) for _ in range(num_blocks)])
        self.conv_out = nn.Conv2d(512, out_channels, 1)

    def forward(self, x):
        x = self.conv_in(x)
        for block in self.blocks:
            x = block(x)
        x = self.conv_out(x)
        return x


class DiffusionFeatureExtractor3(nn.Module):
    def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None, sd: BaseModel = None):
        super().__init__()
        self.version = 3
        self.sd_ref = weakref.ref(sd) if sd is not None else None
        if vae is None:
            vae = AutoencoderTiny.from_pretrained(
                "madebyollin/taef1", torch_dtype=torch.bfloat16)
        self.vae = vae
        # image_encoder_path = "google/siglip-so400m-patch14-384"
        image_encoder_path = "google/siglip2-so400m-patch16-512"
        try:
            self.image_processor = SiglipImageProcessor.from_pretrained(
                image_encoder_path)
        except EnvironmentError:
            self.image_processor = SiglipImageProcessor()
        self.vision_encoder = SiglipVisionModel.from_pretrained(
            image_encoder_path,
            ignore_mismatched_sizes=True
        ).to(device, dtype=dtype)

        self.lpips_model = lpips_model = lpips.LPIPS(net='vgg')
        self.lpips_model = lpips_model.to(device, dtype=torch.float32)
        self.losses = {}
        self.log_every = 100
        self.step = 0
        
    def get_siglip_features(self, tensors_0_1):
        dtype = torch.bfloat16
        device = self.vae.device
        # resize to 384x384
        if 'height' in self.image_processor.size:
            size = self.image_processor.size['height']
        else:
            size = self.image_processor.crop_size['height']
        images = F.interpolate(tensors_0_1, size=(size, size),
                               mode='bicubic', align_corners=False)

        mean = torch.tensor(self.image_processor.image_mean).to(
            device, dtype=dtype
        ).detach()
        std = torch.tensor(self.image_processor.image_std).to(
            device, dtype=dtype
        ).detach()
        # tensors_0_1 = torch.clip((255. * tensors_0_1), 0, 255).round() / 255.0
        clip_image = (
            images - mean.view([1, 3, 1, 1])) / std.view([1, 3, 1, 1])
        id_embeds = self.vision_encoder(
            clip_image,
            output_hidden_states=True,
        )

        last_hidden_state = id_embeds['last_hidden_state']
        return last_hidden_state
    
    def get_lpips_features(self, tensors_0_1):
        device = self.vae.device
        tensors_n1p1 = (tensors_0_1 * 2) - 1
        def get_lpips_features(img):  # -1 to 1
            in0_input = self.lpips_model.scaling_layer(img)
            outs0 = self.lpips_model.net.forward(in0_input)

            feats0 = {}

            feats_list = []
            for kk in range(self.lpips_model.L):
                feats0[kk] = lpips.normalize_tensor(outs0[kk])
                feats_list.append(feats0[kk])

            # 512 in
            # vgg
            # 0 torch.Size([1, 64, 512, 512])
            # 1 torch.Size([1, 128, 256, 256])
            # 2 torch.Size([1, 256, 128, 128])
            # 3 torch.Size([1, 512, 64, 64])
            # 4 torch.Size([1, 512, 32, 32])

            return feats_list

        # do lpips
        lpips_feat_list = [x for x in get_lpips_features(
            tensors_n1p1.to(device, dtype=torch.float32))]
        
        return lpips_feat_list
        

    def forward(
        self,
        noise,
        noise_pred,
        noisy_latents,
        timesteps, 
        batch: DataLoaderBatchDTO, 
        scheduler: CustomFlowMatchEulerDiscreteScheduler,
        # lpips_weight=1.0,
        lpips_weight=10.0,
        clip_weight=0.1,
        pixel_weight=0.1,
        model=None
    ):
        dtype = torch.bfloat16
        device = self.vae.device
        
        
        if model is not None and hasattr(model, 'get_stepped_pred'):
            stepped_latents = model.get_stepped_pred(noise_pred, noise)
        elif self.sd_ref is not None and getattr(self.sd_ref(), "x0_pred", False):
            stepped_latents = noise_pred
        else:
            # stepped_latents = noise - noise_pred
            # first we step the scheduler from current timestep to the very end for a full denoise
            bs = noise_pred.shape[0]
            noise_pred_chunks = torch.chunk(noise_pred, bs)
            timestep_chunks = torch.chunk(timesteps, bs)
            noisy_latent_chunks = torch.chunk(noisy_latents, bs)
            stepped_chunks = []
            for idx in range(bs):
                model_output = noise_pred_chunks[idx]
                timestep = timestep_chunks[idx]
                scheduler._step_index = None
                scheduler._init_step_index(timestep)
                sample = noisy_latent_chunks[idx].to(torch.float32)
                
                sigma = scheduler.sigmas[scheduler.step_index]
                sigma_next = scheduler.sigmas[-1] # use last sigma for final step
                prev_sample = sample + (sigma_next - sigma) * model_output
                stepped_chunks.append(prev_sample)
            
            stepped_latents = torch.cat(stepped_chunks, dim=0)
            
        latents = stepped_latents.to(self.vae.device, dtype=self.vae.dtype)
        
        scaling_factor = self.vae.config.scaling_factor if hasattr(self.vae.config, 'scaling_factor') else 1.0
        shift_factor = self.vae.config.shift_factor if hasattr(self.vae.config, 'shift_factor') else 0.0
        latents = (latents / scaling_factor) + shift_factor
        tensors_n1p1 = self.vae.decode(latents)  # -1 to 1
        if hasattr(tensors_n1p1, 'sample'):
            tensors_n1p1 = tensors_n1p1.sample
        
        pred_images = (tensors_n1p1 + 1) / 2  # 0 to 1
        
        lpips_feat_list_pred = self.get_lpips_features(pred_images.float())
        
        total_loss = 0
        
        with torch.no_grad():
            target_img = batch.tensor.to(device, dtype=dtype)
            # go from -1 to 1 to 0 to 1
            target_img = (target_img + 1) / 2
            lpips_feat_list_target = self.get_lpips_features(target_img.float())
            if clip_weight > 0:
                target_clip_output = self.get_siglip_features(target_img).detach()
        if clip_weight > 0:
            pred_clip_output = self.get_siglip_features(pred_images)
            clip_loss = torch.nn.functional.mse_loss(
                pred_clip_output.float(), target_clip_output.float()
            ) * clip_weight
            
            if 'clip_loss' not in self.losses:
                self.losses['clip_loss'] = clip_loss.item()
            else:
                self.losses['clip_loss'] += clip_loss.item()
            
            total_loss += clip_loss
        
        skip_lpips_layers = []
        
        lpips_loss = 0
        for idx, lpips_feat in enumerate(lpips_feat_list_pred):
            if idx in skip_lpips_layers:
                continue
            lpips_loss += torch.nn.functional.mse_loss(
                lpips_feat.float(), lpips_feat_list_target[idx].float()
            ) * lpips_weight
            
            if f'lpips_loss_{idx}' not in self.losses:
                self.losses[f'lpips_loss_{idx}'] = lpips_loss.item()
            else:
                self.losses[f'lpips_loss_{idx}'] += lpips_loss.item()
            
        total_loss += lpips_loss
            
        # mse_loss = torch.nn.functional.mse_loss(
        #     stepped_latents.float(), batch.latents.float()
        # ) * pixel_weight
        
        # if 'pixel_loss' not in self.losses:
        #     self.losses['pixel_loss'] = mse_loss.item()
        # else:
        #     self.losses['pixel_loss'] += mse_loss.item()
            
        if self.step % self.log_every == 0 and self.step > 0:
            print(f"DFE losses:")
            for key in self.losses:
                self.losses[key] /= self.log_every
                # print in 2.000e-01 format
                print(f" - {key}: {self.losses[key]:.3e}")
            self.losses[key] = 0.0
        
        # total_loss += mse_loss
        self.step += 1
        
        return total_loss

class DiffusionFeatureExtractor4(nn.Module):
    def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None, sd: BaseModel = None):
        super().__init__()
        self.version = 4
        self.sd_ref = weakref.ref(sd) if sd is not None else None
        if vae is None:
            raise ValueError("vae must be provided for DFE4")
        self.vae = vae
        # image_encoder_path = "google/siglip-so400m-patch14-384"
        image_encoder_path = "google/siglip2-so400m-patch16-naflex"
        from transformers import Siglip2ImageProcessor, Siglip2VisionModel
        try:
            self.image_processor = Siglip2ImageProcessor.from_pretrained(
                image_encoder_path)
        except EnvironmentError:
            self.image_processor = Siglip2ImageProcessor()
        
        self.image_processor.max_num_patches = 1024
            
        self.vision_encoder = Siglip2VisionModel.from_pretrained(
            image_encoder_path,
            ignore_mismatched_sizes=True
        ).to(device, dtype=dtype)

        self.losses = {}
        self.log_every = 100
        self.step = 0
    
    def _target_hw(self, h, w, patch, max_patches, eps: float = 1e-5):
        def _snap(x, s):
            x = math.ceil((x * s) / patch) * patch
            return max(patch, int(x))

        lo, hi = eps / 10, 1.0
        while hi - lo >= eps:
            mid = (lo + hi) / 2
            th, tw = _snap(h, mid), _snap(w, mid)
            if (th // patch) * (tw // patch) <= max_patches:
                lo = mid
            else:
                hi = mid
        return _snap(h, lo), _snap(w, lo)


    def tensors_to_siglip_like_features(self, batch: torch.Tensor):
        """
        Args:
            batch: (bs, 3, H, W) tensor already in the desired value range
                   (e.g. [-1, 1] or [0, 1]); no extra rescale / normalize here.

        Returns:
            dict(
                pixel_values        – (bs, L, P) where L = n_h*n_w, P = 3*patch*patch
                pixel_attention_mask– (L,)  all-ones
                spatial_shapes      – (n_h, n_w)
            )
        """
        if batch.ndim != 4:
            raise ValueError("Expected (bs, 3, H, W) tensor")

        bs, c, H, W = batch.shape
        proc        = self.image_processor
        patch       = proc.patch_size
        max_patches = proc.max_num_patches

        # One shared resize for the whole batch
        tgt_h, tgt_w = self._target_hw(H, W, patch, max_patches)
        batch = torch.nn.functional.interpolate(
            batch, size=(tgt_h, tgt_w), mode="bilinear", align_corners=False
        )

        n_h, n_w = tgt_h // patch, tgt_w // patch
        # flat_dim = c * patch * patch
        num_p    = n_h * n_w

        # unfold → (bs, flat_dim, num_p)  → (bs, num_p, flat_dim)
        patches = (
            torch.nn.functional.unfold(batch, kernel_size=patch, stride=patch)
            .transpose(1, 2)
        )

        attn_mask = torch.ones(num_p, dtype=torch.long, device=batch.device)
        spatial   = torch.tensor((n_h, n_w), device=batch.device, dtype=torch.int32)
        
        # repeat attn_mask for each batch element
        attn_mask = attn_mask.unsqueeze(0).repeat(bs, 1)
        spatial = spatial.unsqueeze(0).repeat(bs, 1)

        return {
            "pixel_values": patches,            # (bs, num_patches, patch_dim)
            "pixel_attention_mask": attn_mask,  # (num_patches,)
            "spatial_shapes": spatial
        }
        
    def get_siglip_features(self, tensors_0_1):
        dtype = torch.bfloat16
        device = self.vae.device
        
        tensors_0_1 = torch.clamp(tensors_0_1, 0.0, 1.0)

        mean = torch.tensor(self.image_processor.image_mean).to(
            device, dtype=dtype
        ).detach()
        std = torch.tensor(self.image_processor.image_std).to(
            device, dtype=dtype
        ).detach()
        # tensors_0_1 = torch.clip((255. * tensors_0_1), 0, 255).round() / 255.0
        clip_image = (tensors_0_1 - mean.view([1, 3, 1, 1])) / std.view([1, 3, 1, 1])
        
        encoder_kwargs = self.tensors_to_siglip_like_features(clip_image)
        id_embeds = self.vision_encoder(
            pixel_values=encoder_kwargs['pixel_values'],
            pixel_attention_mask=encoder_kwargs['pixel_attention_mask'],
            spatial_shapes=encoder_kwargs['spatial_shapes'],
            output_hidden_states=True,
        )

        image_embeds = id_embeds['hidden_states'][-2]  # penultimate layer
        # image_embeds = id_embeds['pooler_output']
        # image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
        return image_embeds
    
    def step_latents(self, noise, noise_pred, noisy_latents, timesteps, scheduler):
        bs = noise_pred.shape[0]
        noise_pred_chunks = torch.chunk(noise_pred, bs)
        timestep_chunks = torch.chunk(timesteps, bs)
        noisy_latent_chunks = torch.chunk(noisy_latents, bs)
        stepped_chunks = []
        for idx in range(bs):
            model_output = noise_pred_chunks[idx]
            timestep = timestep_chunks[idx]
            scheduler._step_index = None
            scheduler._init_step_index(timestep)
            sample = noisy_latent_chunks[idx].to(torch.float32)
            
            sigma = scheduler.sigmas[scheduler.step_index]
            sigma_next = scheduler.sigmas[-1] # use last sigma for final step
            prev_sample = sample + (sigma_next - sigma) * model_output
            stepped_chunks.append(prev_sample)
        
        stepped_latents = torch.cat(stepped_chunks, dim=0)
        return stepped_latents

    def forward(
        self,
        noise,
        noise_pred,
        noisy_latents,
        timesteps, 
        batch: DataLoaderBatchDTO, 
        scheduler: CustomFlowMatchEulerDiscreteScheduler,
        clip_weight=1.0,
        mse_weight=0.0,
        model=None
    ):
        dtype = torch.bfloat16
        device = self.vae.device
        tensors = batch.tensor.to(device, dtype=dtype)
        is_video = False
        if len(noise_pred.shape) == 5:
            # (B, C, T, H, W): video VAEs decode whole clips (latent frames do
            # not map 1:1 to pixel frames), so the latents stay 5D through the
            # VAE and the decoded PIXEL frames fold into the batch dim below,
            # matching the folded target frames
            is_video = True
        
        if len(tensors.shape) == 5:
            # batch tensor is frames-first (B, T, C, H, W): fold to (B*T, C, H, W)
            tensors = tensors.reshape(-1, *tensors.shape[2:])
            
        if model is not None and hasattr(model, 'get_stepped_pred'):
            stepped_latents = model.get_stepped_pred(noise_pred, noise)
        else:
            with torch.no_grad():
                tv = timesteps.to(noise_pred.device).to(noise_pred.dtype) / 1000.0
                # expand shape to match noise_pred
                while len(tv.shape) < len(noise_pred.shape):
                    tv = tv.unsqueeze(-1)
                    # min 0.001
                    tv = torch.clamp(tv, min=0.001)

            if self.sd_ref is not None and getattr(self.sd_ref(), "get_x0_prediction", None) is not None:
                x0 = self.sd_ref().get_x0_prediction(noise_pred, noisy_latents, tv)
            elif self.sd_ref is not None and getattr(self.sd_ref(), "x0_pred", False):
                x0 = noise_pred
            else:
                # step latent
                x0 = noisy_latents - tv * noise_pred

            stepped_latents = x0

        latents = stepped_latents.to(self.vae.device, dtype=self.vae.dtype)

        scaling_factor = self.vae.config.scaling_factor if hasattr(self.vae.config, 'scaling_factor') else 1.0
        shift_factor = self.vae.config.shift_factor if hasattr(self.vae.config, 'shift_factor') else 0.0
        latents = (latents / scaling_factor) + shift_factor
        # video latents stay 5D (B, C, T, H, W): the video VAE decodes the whole
        # clip, and the decoded PIXEL frames fold into the batch dim to match
        # the folded target frames
        tensors_n1p1 = self.vae.decode(latents)  # -1 to 1
        if hasattr(tensors_n1p1, 'sample'):
            tensors_n1p1 = tensors_n1p1.sample
        
        if is_video:
            tensors_n1p1 = _fold_frames_to_batch(tensors_n1p1)
        
        pred_images = (tensors_n1p1 + 1) / 2  # 0 to 1
        
        total_loss = 0
        
        with torch.no_grad():
            target_img = tensors.to(device, dtype=dtype)
            # go from -1 to 1 to 0 to 1
            target_img = (target_img + 1) / 2
            if clip_weight > 0:
                target_clip_output = self.get_siglip_features(target_img).detach()
        if clip_weight > 0:
            pred_clip_output = self.get_siglip_features(pred_images)
            clip_loss = torch.nn.functional.mse_loss(
                pred_clip_output.float(), target_clip_output.float()
            ) * clip_weight
            
            if 'clip_loss' not in self.losses:
                self.losses['clip_loss'] = clip_loss.item()
            else:
                self.losses['clip_loss'] += clip_loss.item()
            
            total_loss += clip_loss
        if mse_weight > 0:
            mse_loss = torch.nn.functional.mse_loss(
                pred_images.float(), target_img.float()
            ) * mse_weight
            
            if 'mse_loss' not in self.losses:
                self.losses['mse_loss'] = mse_loss.item()
            else:
                self.losses['mse_loss'] += mse_loss.item()
            
            total_loss += mse_loss

        if self.step % self.log_every == 0 and self.step > 0:
            print(f"DFE losses:")
            for key in self.losses:
                self.losses[key] /= self.log_every
                # print in 2.000e-01 format
                print(f" - {key}: {self.losses[key]:.3e}")
                self.losses[key] = 0.0
        
        # total_loss += mse_loss
        self.step += 1
        
        return total_loss
    
class DiffusionFeatureExtractor5(DiffusionFeatureExtractor4):
    def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None, sd: BaseModel = None):
        super().__init__(device=device, dtype=dtype, vae=vae, sd=sd)
        self.version = 5
    
    def step_latents(self, noise, noise_pred, noisy_latents, timesteps, scheduler, total_steps: int = 1000, eps: float = 1e-6):
        bs = noise_pred.shape[0]

        # Chunk inputs per-sample (keeps existing structure)
        noise_pred_chunks   = torch.chunk(noise_pred,   bs)
        timestep_chunks     = torch.chunk(timesteps,    bs)
        noisy_latent_chunks = torch.chunk(noisy_latents, bs)
        noise_chunks        = torch.chunk(noise,        bs)

        stepped_chunks = []
        x0_pred_chunks = []

        for idx in range(bs):
            model_output = noise_pred_chunks[idx]              # predicted noise (same shape as latent)
            timestep     = timestep_chunks[idx]                # scalar tensor per sample (e.g., [t])
            sample       = noisy_latent_chunks[idx].to(torch.float32)
            noise_i      = noise_chunks[idx].to(sample.dtype).to(sample.device)

            # Initialize scheduler step index for this sample
            scheduler._step_index = None
            scheduler._init_step_index(timestep)

            # ---- Step +50 indices (or to the end) in sigma-space ----
            sigma      = scheduler.sigmas[scheduler.step_index]
            target_idx = min(scheduler.step_index + 50, len(scheduler.sigmas) - 1)
            sigma_next = scheduler.sigmas[target_idx]

            # One-step update along the model-predicted direction
            stepped = sample + (sigma_next - sigma) * model_output
            stepped_chunks.append(stepped)

            # ---- Inverse-Gaussian recovery at the target timestep ----
            t_01 = (scheduler.sigmas[target_idx]).to(stepped.device).to(stepped.dtype)
            original_samples = (stepped - t_01 * noise_i) / (1.0 - t_01)
            x0_pred_chunks.append(original_samples)

        # stepped_latents   = torch.cat(stepped_chunks, dim=0)
        predicted_images  = torch.cat(x0_pred_chunks, dim=0)
        # return stepped_latents, predicted_images
        return predicted_images


class DiffusionFeatureExtractor6(nn.Module):
    def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None, sd: BaseModel = None):
        super().__init__()
        self.version = 6
        self.sd_ref = weakref.ref(sd) if sd is not None else None
        if vae is None:
            raise ValueError("vae must be provided for DFE4")
        self.vae = vae
        # pretrained_model_name = "facebook/dinov3-vits16-pretrain-lvd1689m"
        # pretrained_model_name = "facebook/dinov3-vitl16-pretrain-lvd1689m"
        pretrained_model_name = "facebook/dinov3-vith16plus-pretrain-lvd1689m"
        # pretrained_model_name = "facebook/dinov3-vit7b16-pretrain-lvd1689m"
        self.processor = AutoImageProcessor.from_pretrained(pretrained_model_name)
        self.model = AutoModel.from_pretrained(
            pretrained_model_name, 
            device_map=device,
            dtype=dtype,
        ).to(device, dtype=dtype)

        self.losses = {}
        self.log_every = 100
        self.step = 0

        # cache normalization constants once so prepare_inputs doesn't rebuild
        # them (CPU tensor construction + H2D copy) on every call
        self.image_mean = torch.tensor(self.processor.image_mean, device=device, dtype=dtype).view(1, 3, 1, 1)
        self.image_std = torch.tensor(self.processor.image_std, device=device, dtype=dtype).view(1, 3, 1, 1)

        # last_hidden_state token order: [CLS, registers..., patches]
        self.num_prefix_tokens = 1 + getattr(self.model.config, "num_register_tokens", 0)
        # CLS is a global anchor only; patch tokens carry the spatial signal
        self.cls_weight = 0.1

    def _dino_features(self, inputs):
        hidden = self.model(**inputs).last_hidden_state
        return hidden[:, 0], hidden[:, self.num_prefix_tokens:]

    def prepare_inputs(self, tensor_0_1: torch.Tensor):
        """
        tensor_0_1: (bs, 3, h, w), float, values in [0, 1]
        returns: {"pixel_values": (bs, 3, H, W)} ready for the vision transformer
        """

        if tensor_0_1.ndim != 4 or tensor_0_1.shape[1] != 3:
            raise ValueError(f"Expected (bs, 3, h, w), got {tuple(tensor_0_1.shape)}")

        x = tensor_0_1
        if not torch.is_floating_point(x):
            x = x.float()

        # VAE decode can overshoot [0, 1] slightly; clamp back into range.
        # Inputs are documented as 0..1 and clamped here, so no 0..255 rescale
        # heuristic is needed (the old max().item() check forced a GPU sync
        # every call and could never trigger after this clamp anyway).
        x = torch.clamp(x, 0.0, 1.0)

        # Resize
        # if not divisible by 16 or total pixels > max_res*max_res, resize to fit within 16 patches
        max_res = 512
        p = 16
        if (x.shape[-1] % p != 0) or (x.shape[-2] % p != 0) or (x.shape[-1] * x.shape[-2] > max_res * max_res):
            target_h = x.shape[-2]
            target_w = x.shape[-1]
            if x.shape[-1] * target_h > max_res * max_res:
                scale_factor = math.sqrt((max_res * max_res) / (target_w * target_h))
                target_h = int(target_h * scale_factor)
                target_w = int(target_w * scale_factor)
            target_h = (target_h // p) * p
            target_w = (target_w // p) * p
            x = F.interpolate(x, size=(target_h, target_w), mode="bilinear", align_corners=False)
            
        # Normalize
        if self.processor.do_normalize:
            mean = self.image_mean.to(device=x.device, dtype=x.dtype)
            std = self.image_std.to(device=x.device, dtype=x.dtype)
            x = (x - mean) / std

        return {"pixel_values": x}
    
    def forward(
        self,
        noise,
        noise_pred,
        noisy_latents,
        timesteps, 
        batch: DataLoaderBatchDTO, 
        scheduler: CustomFlowMatchEulerDiscreteScheduler,
        model=None
    ):
        dtype = torch.bfloat16
        device = self.vae.device
        tensors = batch.tensor.to(device, dtype=dtype)
        is_video = False
        if len(noise_pred.shape) == 5:
            # (B, C, T, H, W): video VAEs decode whole clips (latent frames do
            # not map 1:1 to pixel frames), so the latents stay 5D through the
            # VAE and the decoded PIXEL frames fold into the batch dim below,
            # matching the folded target frames
            is_video = True
        
        if len(tensors.shape) == 5:
            # batch tensor is frames-first (B, T, C, H, W): fold to (B*T, C, H, W)
            tensors = tensors.reshape(-1, *tensors.shape[2:])
            
        with torch.no_grad():
            tv = timesteps.to(noise_pred.device).to(noise_pred.dtype) / 1000.0
            # expand shape to match noise_pred
            while len(tv.shape) < len(noise_pred.shape):
                tv = tv.unsqueeze(-1)
                # min 0.001
                tv = torch.clamp(tv, min=0.001)
            
        if self.sd_ref is not None and getattr(self.sd_ref(), "get_x0_prediction", None) is not None:
            x0 = self.sd_ref().get_x0_prediction(noise_pred, noisy_latents, tv)
        elif self.sd_ref is not None and getattr(self.sd_ref(), "x0_pred", False):
            x0 = noise_pred
        else:
            # step latent
            x0 = noisy_latents - tv * noise_pred
        
        stepped_latents = x0
            
        latents = stepped_latents.to(self.vae.device, dtype=self.vae.dtype)

        scaling_factor = self.vae.config.scaling_factor if hasattr(self.vae.config, 'scaling_factor') else 1.0
        shift_factor = self.vae.config.shift_factor if hasattr(self.vae.config, 'shift_factor') else 0.0
        latents = (latents / scaling_factor) + shift_factor
        # video latents stay 5D (B, C, T, H, W): the video VAE decodes the whole
        # clip, and the decoded PIXEL frames fold into the batch dim to match
        # the folded target frames
        tensors_n1p1 = self.vae.decode(latents)  # -1 to 1
        if hasattr(tensors_n1p1, 'sample'):
            tensors_n1p1 = tensors_n1p1.sample
        
        if is_video:
            tensors_n1p1 = _fold_frames_to_batch(tensors_n1p1)
        
        pred_images = (tensors_n1p1 + 1) / 2  # 0 to 1

        with torch.no_grad():
            target_img = tensors.to(device, dtype=dtype)
            # go from -1 to 1 to 0 to 1
            target_img = (target_img + 1) / 2
            target_dino_input = self.prepare_inputs(target_img)
            target_cls, target_patches = self._dino_features(target_dino_input)
            target_cls = target_cls.detach()
            target_patches = target_patches.detach()
        pred_dino_input = self.prepare_inputs(pred_images)
        pred_cls, pred_patches = self._dino_features(pred_dino_input)

        # per-token cosine so no single bright region dominates
        patch_loss = (1.0 - F.cosine_similarity(pred_patches.float(), target_patches.float(), dim=-1)).mean(dim=1)
        cls_loss = 1.0 - F.cosine_similarity(pred_cls.float(), target_cls.float(), dim=-1)
        # (B*T,) -> (B,) so video frames fold back into their source sample
        bs = noise_pred.shape[0]
        patch_loss = patch_loss.view(bs, -1).mean(dim=1)
        cls_loss = cls_loss.view(bs, -1).mean(dim=1)
        # per-sample (B,); trainer applies the model's per-sample loss scaling and reduces
        dino_loss = patch_loss + self.cls_weight * cls_loss
        
        # accumulate on-GPU; .item() every step forces a full pipeline sync,
        # so only sync when we actually log
        for key, val in (('dinov3_patch', patch_loss.mean()), ('dinov3_cls', cls_loss.mean())):
            if key not in self.losses:
                self.losses[key] = val.detach()
            else:
                self.losses[key] = self.losses[key] + val.detach()

        with torch.no_grad():
            if self.step % self.log_every == 0 and self.step > 0:
                print(f"DFE losses:")
                for key in self.losses:
                    avg = self.losses[key] / self.log_every
                    if torch.is_tensor(avg):
                        avg = avg.item()
                    # print in 2.000e-01 format
                    print(f" - {key}: {avg:.3e}")
                    self.losses[key] = 0.0
            
            # total_loss += mse_loss
            self.step += 1
        
        return dino_loss
    
class ModelOutputWrapper:
    def __init__(self, head, depth, normals, segmentation):
        self.head = head
        self.depth = depth
        self.normals = normals
        self.segmentation = segmentation
    

class DiffusionFeatureExtractor7(nn.Module):
    def __init__(
        self, 
        device=torch.device("cuda"), 
        dtype=torch.bfloat16, 
        vae=None, 
        sd=None,
        partial_step: bool = False
    ):
        super().__init__()
        
        self.version = 7
        self.sd_ref = weakref.ref(sd) if sd is not None else None
        from toolkit.models.tipsv2 import TIPSv2DPTModel
        pretrained_model_name = "google/tipsv2-b14-dpt"
        self.model = TIPSv2DPTModel.from_pretrained(
            pretrained_model_name,
            dtype=dtype,
        ).to(device, dtype=dtype)

        self.losses = {}
        self.log_every = 100
        self.step = 0
        self.do_partial_step = partial_step
    
    def get_pred(self, tensor_0_1: torch.Tensor):
        """
        tensor_0_1: (bs, 3, h, w), float, values in [0, 1]
        returns: {"pixel_values": (bs, 3, H, W)} ready for the vision transformer
        """

        if tensor_0_1.ndim != 4 or tensor_0_1.shape[1] != 3:
            raise ValueError(f"Expected (bs, 3, h, w), got {tuple(tensor_0_1.shape)}")

        x = tensor_0_1.to(self.model.device, dtype=self.model.dtype)

        # Resize
        # if not divisible by 16 or total pixels > max_res*max_res, resize to fit within 16 patches
        max_res = 1024
        p = 14
        if (x.shape[-1] % p != 0) or (x.shape[-2] % p != 0) or (x.shape[-1] * x.shape[-2] > max_res * max_res):
            target_h = x.shape[-2]
            target_w = x.shape[-1]
            if x.shape[-1] * target_h > max_res * max_res:
                scale_factor = math.sqrt((max_res * max_res) / (target_w * target_h))
                target_h = int(target_h * scale_factor)
                target_w = int(target_w * scale_factor)
            target_h = (target_h // p) * p
            target_w = (target_w // p) * p
            x = F.interpolate(x, size=(target_h, target_w), mode="bilinear", align_corners=False)
        
        # do inference. us standard dpy but also the head
        pixel_values = x.to(self.model.device, dtype=self.model.dtype)
        h, w = pixel_values.shape[2:]
        dpt_inputs = self.model._extract_intermediate(pixel_values)
        # head is a list of 4
        # each of the 4 is a tuple of (embeds, hidden_state)
        # concat the hidden states from the 4 layers on dim 1
        head = torch.cat([h[1] for h in dpt_inputs], dim=1)
        return ModelOutputWrapper(
            head=head,
            depth=self.model.depth_head(dpt_inputs, image_size=(h, w)),
            normals=self.model.normals_head(dpt_inputs, image_size=(h, w)),
            segmentation=self.model.segmentation_head(dpt_inputs, image_size=(h, w)),
        )
    
    def forward(
        self,
        noise,
        noise_pred,
        noisy_latents,
        timesteps, 
        batch: DataLoaderBatchDTO, 
        scheduler: CustomFlowMatchEulerDiscreteScheduler,
        model=None
    ):
        dtype = torch.bfloat16
        device = self.sd_ref().vae.device
        tensors = batch.tensor.to(device, dtype=dtype)
        is_video = False
        if len(noise_pred.shape) == 5:
            # (B, C, T, H, W): video VAEs decode whole clips (latent frames do
            # not map 1:1 to pixel frames), so the latents stay 5D through the
            # VAE and the decoded PIXEL frames fold into the batch dim below,
            # matching the folded target frames
            is_video = True
        
        if len(tensors.shape) == 5:
            # batch tensor is frames-first (B, T, C, H, W): fold to (B*T, C, H, W)
            tensors = tensors.reshape(-1, *tensors.shape[2:])
            
        with torch.no_grad():
            tv = timesteps.to(noise_pred.device).to(noise_pred.dtype) / 1000.0
            # expand shape to match noise_pred
            while len(tv.shape) < len(noise_pred.shape):
                tv = tv.unsqueeze(-1)
        
        with torch.no_grad():
            target_0_1 = (tensors + 1) / 2  # 0 to 1
        
        if not self.do_partial_step:
            if getattr(self.sd_ref(), "get_x0_prediction", None) is not None:
                x0 = self.sd_ref().get_x0_prediction(noise_pred, noisy_latents, tv)
            elif getattr(self.sd_ref(), "x0_pred", False):
                x0 = noise_pred
            else:
                # step latent
                x0 = noisy_latents - tv * noise_pred
            stepped_latents = x0
            # min 0.001
            tv = torch.clamp(tv, min=0.001)
        else:
            # step is random 0.05 to 0.02
            step = torch.rand_like(tv) * 0.03 + 0.02
            next_step = tv - step
            next_step = torch.clamp(next_step, min=0.0)
            stepped_latents = noisy_latents + (next_step - tv) * noise_pred
            
            with torch.no_grad():
                # make a noisy target at next timestep
                target_latents = batch.latents.to(self.sd_ref().vae.device, dtype=self.sd_ref().vae.dtype)
                # add noise
                target_latents = (1.0 - next_step) * target_latents + next_step * noise
                target_n1p1 = self.sd_ref().decode_latents(target_latents)
                if is_video:
                    target_n1p1 = _fold_frames_to_batch(target_n1p1)
                target_0_1 = (target_n1p1 + 1) / 2  # 0 to 1

        latents = stepped_latents.to(self.sd_ref().vae.device, dtype=self.sd_ref().vae.dtype)
        
        tensors_n1p1 = self.sd_ref().decode_latents(latents)
        if is_video:
            # (B, 3, T_px, H, W) -> (B*T_px, 3, H, W), and repeat each sample's
            # timestep weight for every one of its decoded frames
            b, t_px = tensors_n1p1.shape[0], tensors_n1p1.shape[2]
            tensors_n1p1 = _fold_frames_to_batch(tensors_n1p1)
            tv = tv.reshape(b, 1, 1, 1).repeat_interleave(t_px, dim=0)
        
        pred_images = (tensors_n1p1 + 1) / 2  # 0 to 1
        
        device = self.model.device
        dtype = self.model.dtype
        
        # velocity_equiv_weight = (1.0 / torch.clamp(tv, min=0.1) ** 2)
        velocity_equiv_weight = 1.0

        with torch.no_grad():
            target = self.get_pred(target_0_1)
        
        pred_images = pred_images.to(device, dtype=dtype)
        pred = self.get_pred(pred_images)
        
        head_loss = torch.nn.functional.mse_loss(
            pred.head.float(), target.head.float(), reduction='none'
        ) * velocity_equiv_weight
        head_loss = head_loss.mean()
        
        depth_loss = torch.nn.functional.l1_loss(
            pred.depth.float(), target.depth.float(), reduction='none'
        ) * velocity_equiv_weight
        depth_loss = depth_loss.mean()
        
        normals_loss = torch.nn.functional.l1_loss(
            pred.normals.float(), target.normals.float(), reduction='none'
        ) * velocity_equiv_weight
        normals_loss = normals_loss.mean()
        
        segmentation_loss = torch.nn.functional.l1_loss(
            pred.segmentation.float(), target.segmentation.float(), reduction='none'
        ) * velocity_equiv_weight
        segmentation_loss = segmentation_loss.mean()
        
        total_loss = (head_loss + depth_loss + normals_loss + segmentation_loss) / 4.0
        
        if self.do_partial_step:
            total_loss = total_loss * 10.0
        
        if 'total' not in self.losses:
            self.losses['total'] = total_loss.item()
        else:
            self.losses['total'] += total_loss.item()
            
        if 'head' not in self.losses:
            self.losses['head'] = head_loss.item()
        else:
            self.losses['head'] += head_loss.item()
            
        if 'depth' not in self.losses:
            self.losses['depth'] = depth_loss.item()
        else:
            self.losses['depth'] += depth_loss.item()
            
        if 'normals' not in self.losses:
            self.losses['normals'] = normals_loss.item()
        else:
            self.losses['normals'] += normals_loss.item()
            
        if 'segmentation' not in self.losses:
            self.losses['segmentation'] = segmentation_loss.item()
        else:   
            self.losses['segmentation'] += segmentation_loss.item()
        
        with torch.no_grad():
            if self.step % self.log_every == 0 and self.step > 0:
                print(f"DFE losses:")
                for key in self.losses:
                    self.losses[key] /= self.log_every
                    # print in 2.000e-01 format
                    print(f" - {key}: {self.losses[key]:.3e}")
                self.losses[key] = 0.0
            
            # total_loss += mse_loss
            self.step += 1
        
        return total_loss

class DiffusionFeatureExtractor8(DiffusionFeatureExtractor7):
    def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None, sd=None):
        super().__init__(device=device, dtype=dtype, vae=vae, sd=sd, partial_step=True)
        self.version = 8

class DiffusionFeatureExtractor9(nn.Module):
    def __init__(
        self, 
        device=torch.device("cuda"), 
        dtype=torch.bfloat16, 
        vae=None, 
        sd=None,
        partial_step: bool = False
    ):
        super().__init__()
        
        self.version = 9
        self.sd_ref = weakref.ref(sd) if sd is not None else None
        ckpt_path = huggingface_hub.hf_hub_download(repo_id="facebook/sapiens2-pretrain-1b", filename="sapiens2_1b_pretrain.safetensors")
        self.model = Sapiens2(arch="sapiens2_1b", img_size=(1024, 768), patch_size=16).eval().cuda()  # img_size is (H, W)
        self.model.load_state_dict(load_file(ckpt_path))
        self.model.to(device, dtype=dtype)

        self.losses = {}
        self.log_every = 100
        self.step = 0
        self.do_partial_step = partial_step
    
    def get_pred(self, tensor_0_1: torch.Tensor):
        if tensor_0_1.ndim != 4 or tensor_0_1.shape[1] != 3:
            raise ValueError(f"Expected (bs, 3, h, w), got {tuple(tensor_0_1.shape)}")

        x = tensor_0_1.to(self.model.device, dtype=self.model.dtype)
        """Apply ImageNet normalization to a (B, C, H, W) RGB tensor in [0, 1]."""
        mean = torch.as_tensor((0.485, 0.456, 0.406), dtype=x.dtype, device=x.device).view(1, 3, 1, 1)
        std = torch.as_tensor((0.229, 0.224, 0.225), dtype=x.dtype, device=x.device).view(1, 3, 1, 1)
        x = (x - mean) / std

        # Resize
        # if not divisible by 16 or total pixels > max_res*max_res, resize to fit within 16 patches
        max_res = 1024
        p = 16
        if (x.shape[-1] % p != 0) or (x.shape[-2] % p != 0) or (x.shape[-1] * x.shape[-2] > max_res * max_res):
            target_h = x.shape[-2]
            target_w = x.shape[-1]
            if x.shape[-1] * target_h > max_res * max_res:
                scale_factor = math.sqrt((max_res * max_res) / (target_w * target_h))
                target_h = int(target_h * scale_factor)
                target_w = int(target_w * scale_factor)
            target_h = (target_h // p) * p
            target_w = (target_w // p) * p
            x = F.interpolate(x, size=(target_h, target_w), mode="bilinear", align_corners=False)
        x = x.to(self.model.device, dtype=self.model.dtype)
        features = self.model(x)[0]
        return features
    
    def forward(
        self,
        noise,
        noise_pred,
        noisy_latents,
        timesteps, 
        batch: DataLoaderBatchDTO, 
        scheduler: CustomFlowMatchEulerDiscreteScheduler,
        model=None
    ):
        dtype = torch.bfloat16
        device = self.sd_ref().vae.device
        tensors = batch.tensor.to(device, dtype=dtype)
        is_video = False
        if len(noise_pred.shape) == 5:
            # (B, C, T, H, W): video VAEs decode whole clips (latent frames do
            # not map 1:1 to pixel frames), so the latents stay 5D through the
            # VAE and the decoded PIXEL frames fold into the batch dim below,
            # matching the folded target frames
            is_video = True
        
        if len(tensors.shape) == 5:
            # batch tensor is frames-first (B, T, C, H, W): fold to (B*T, C, H, W)
            tensors = tensors.reshape(-1, *tensors.shape[2:])
            
        with torch.no_grad():
            tv = timesteps.to(noise_pred.device).to(noise_pred.dtype) / 1000.0
            # expand shape to match noise_pred
            while len(tv.shape) < len(noise_pred.shape):
                tv = tv.unsqueeze(-1)
        
        with torch.no_grad():
            target_0_1 = (tensors + 1) / 2  # 0 to 1
        
        if not self.do_partial_step:
            if getattr(self.sd_ref(), "get_x0_prediction", None) is not None:
                x0 = self.sd_ref().get_x0_prediction(noise_pred, noisy_latents, tv)
            elif getattr(self.sd_ref(), "x0_pred", False):
                x0 = noise_pred
            else:
                # step latent
                x0 = noisy_latents - tv * noise_pred
            stepped_latents = x0
            # min 0.001
            tv = torch.clamp(tv, min=0.001)
        else:
            # step is random 0.1 to 0.25
            step = torch.rand_like(tv) * 0.15 + 0.1
            next_step = tv - step
            next_step = torch.clamp(next_step, min=0.0)
            stepped_latents = noisy_latents + (next_step - tv) * noise_pred

            with torch.no_grad():
                # make a noisy target at next timestep
                target_latents = batch.latents.to(self.sd_ref().vae.device, dtype=self.sd_ref().vae.dtype)
                # add noise
                target_latents = (1.0 - next_step) * target_latents + next_step * noise
                target_n1p1 = self.sd_ref().decode_latents(target_latents)
                if is_video:
                    target_n1p1 = _fold_frames_to_batch(target_n1p1)
                target_0_1 = (target_n1p1 + 1) / 2  # 0 to 1

        latents = stepped_latents.to(self.sd_ref().vae.device, dtype=self.sd_ref().vae.dtype)
        
        tensors_n1p1 = self.sd_ref().decode_latents(latents)
        if is_video:
            # (B, 3, T_px, H, W) -> (B*T_px, 3, H, W), and repeat each sample's
            # timestep weight for every one of its decoded frames
            b, t_px = tensors_n1p1.shape[0], tensors_n1p1.shape[2]
            tensors_n1p1 = _fold_frames_to_batch(tensors_n1p1)
            tv = tv.reshape(b, 1, 1, 1).repeat_interleave(t_px, dim=0)
        
        pred_images = (tensors_n1p1 + 1) / 2  # 0 to 1
        
        device = self.model.device
        dtype = self.model.dtype

        with torch.no_grad():
            target = self.get_pred(target_0_1)
        
        pred_images = pred_images.to(device, dtype=dtype)
        pred = self.get_pred(pred_images)
        
        perceptual_loss = torch.nn.functional.mse_loss(
            pred.float(), target.float(), reduction="none"
        )
        velocity_equiv_weight = (1.0 / torch.clamp(tv, min=0.01) ** 2)
        loss_perceptual = (perceptual_loss * velocity_equiv_weight).mean()
        
        if self.do_partial_step:
            loss_perceptual = loss_perceptual * 10.0
        
        if 'loss' not in self.losses:
            self.losses['loss'] = loss_perceptual.item()
        else:
            self.losses['loss'] += loss_perceptual.item()
        with torch.no_grad():
            if self.step % self.log_every == 0 and self.step > 0:
                print(f"DFE losses:")
                for key in self.losses:
                    self.losses[key] /= self.log_every
                    # print in 2.000e-01 format
                    print(f" - {key}: {self.losses[key]:.3e}")
                self.losses[key] = 0.0
            
            # total_loss += mse_loss
            self.step += 1
        
        return loss_perceptual

class DiffusionFeatureExtractor10(nn.Module):
    def __init__(
        self,
        device=torch.device("cuda"),
        dtype=torch.bfloat16,
        vae=None,
        sd=None,
        partial_step: bool = False
    ):
        super().__init__()

        self.version = 10
        self.sd_ref = weakref.ref(sd) if sd is not None else None
        self.lpips_model = lpips.LPIPS(net='vgg')
        self.lpips_model = self.lpips_model.to(device, dtype=torch.float32)

        self.losses = {}
        self.log_every = 100
        self.step = 0
        self.do_partial_step = partial_step

    def _vgg_slices(self, x):
        # run the lpips vgg backbone slice-by-slice so we can gradient
        # checkpoint each slice. checkpointing activates whenever grads are
        # enabled, so it does not require the module to be in train mode.
        net = self.lpips_model.net
        slices = [net.slice1, net.slice2, net.slice3, net.slice4, net.slice5]
        outs = []
        h = x
        for s in slices:
            if torch.is_grad_enabled():
                h = ckpt.checkpoint(s, h, use_reentrant=False)
            else:
                h = s(h)
            outs.append(h)
        return outs

    def get_lpips_features(self, tensors_0_1):
        device = self.lpips_model.scaling_layer.shift.device
        tensors_n1p1 = (tensors_0_1 * 2) - 1
        def get_lpips_features(img):  # -1 to 1
            in0_input = self.lpips_model.scaling_layer(img)
            outs0 = self._vgg_slices(in0_input)

            feats_list = []
            for kk in range(self.lpips_model.L):
                feats_list.append(lpips.normalize_tensor(outs0[kk]))

            return feats_list

        lpips_feat_list = [x for x in get_lpips_features(
            tensors_n1p1.to(device, dtype=torch.float32))]

        return lpips_feat_list

    def forward(
        self,
        noise,
        noise_pred,
        noisy_latents,
        timesteps,
        batch: DataLoaderBatchDTO,
        scheduler: CustomFlowMatchEulerDiscreteScheduler,
        model=None
    ):
        dtype = torch.bfloat16
        device = self.sd_ref().vae.device
        tensors = batch.tensor.to(device, dtype=dtype)
        is_video = False
        if len(noise_pred.shape) == 5:
            # (B, C, T, H, W): video VAEs decode whole clips (latent frames do
            # not map 1:1 to pixel frames), so the latents stay 5D through the
            # VAE and the decoded PIXEL frames fold into the batch dim below,
            # matching the folded target frames
            is_video = True

        if len(tensors.shape) == 5:
            # batch tensor is frames-first (B, T, C, H, W): fold to (B*T, C, H, W)
            tensors = tensors.reshape(-1, *tensors.shape[2:])

        with torch.no_grad():
            tv = timesteps.to(noise_pred.device).to(noise_pred.dtype) / 1000.0
            # expand shape to match noise_pred
            while len(tv.shape) < len(noise_pred.shape):
                tv = tv.unsqueeze(-1)

        with torch.no_grad():
            target_0_1 = (tensors + 1) / 2  # 0 to 1

        if not self.do_partial_step:
            if self.sd_ref is not None and getattr(self.sd_ref(), "get_x0_prediction", None) is not None:
                x0 = self.sd_ref().get_x0_prediction(noise_pred, noisy_latents, tv)
            elif self.sd_ref is not None and getattr(self.sd_ref(), "x0_pred", False):
                x0 = noise_pred
            else:
                # step latent
                x0 = noisy_latents - tv * noise_pred
            stepped_latents = x0
            # min 0.001
            tv = torch.clamp(tv, min=0.001)
        else:
            # step is random 0.1 to 0.25
            step = torch.rand_like(tv) * 0.15 + 0.1
            next_step = tv - step
            next_step = torch.clamp(next_step, min=0.0)
            stepped_latents = noisy_latents + (next_step - tv) * noise_pred

            with torch.no_grad():
                # make a noisy target at next timestep
                target_latents = batch.latents.to(self.sd_ref().vae.device, dtype=self.sd_ref().vae.dtype)
                # add noise
                target_latents = (1.0 - next_step) * target_latents + next_step * noise
                target_n1p1 = self.sd_ref().decode_latents(target_latents)
                if is_video:
                    target_n1p1 = _fold_frames_to_batch(target_n1p1)
                target_0_1 = (target_n1p1 + 1) / 2  # 0 to 1

        latents = stepped_latents.to(self.sd_ref().vae.device, dtype=self.sd_ref().vae.dtype)

        tensors_n1p1 = self.sd_ref().decode_latents(latents)
        if is_video:
            # (B, 3, T_px, H, W) -> (B*T_px, 3, H, W), and repeat each sample's
            # timestep weight for every one of its decoded frames
            b, t_px = tensors_n1p1.shape[0], tensors_n1p1.shape[2]
            tensors_n1p1 = _fold_frames_to_batch(tensors_n1p1)
            tv = tv.reshape(b, 1, 1, 1).repeat_interleave(t_px, dim=0)

        pred_images = (tensors_n1p1 + 1) / 2  # 0 to 1

        with torch.no_grad():
            target_feats = self.get_lpips_features(target_0_1.float())

        pred_feats = self.get_lpips_features(pred_images.float())

        velocity_equiv_weight = (1.0 / torch.clamp(tv, min=0.01) ** 2)

        loss_perceptual = 0
        for idx, pred_feat in enumerate(pred_feats):
            perceptual_loss = torch.nn.functional.mse_loss(
                pred_feat.float(), target_feats[idx].float(), reduction="none"
            )
            # mean over channels/spatial per sample, keep batch dim to weight by timestep
            perceptual_loss = perceptual_loss.mean(dim=[1, 2, 3], keepdim=True)
            loss_perceptual = loss_perceptual + (perceptual_loss * velocity_equiv_weight).mean()

        if self.do_partial_step:
            loss_perceptual = loss_perceptual * 10.0

        if 'loss' not in self.losses:
            self.losses['loss'] = loss_perceptual.item()
        else:
            self.losses['loss'] += loss_perceptual.item()
        with torch.no_grad():
            if self.step % self.log_every == 0 and self.step > 0:
                print(f"DFE losses:")
                for key in self.losses:
                    self.losses[key] /= self.log_every
                    # print in 2.000e-01 format
                    print(f" - {key}: {self.losses[key]:.3e}")
                self.losses[key] = 0.0

            # total_loss += mse_loss
            self.step += 1

        return loss_perceptual

def load_dfe(model_path, vae=None, sd: 'BaseModel' = None) -> DiffusionFeatureExtractor:
    if model_path == "v3":
        dfe = DiffusionFeatureExtractor3(vae=vae, sd=sd)
        dfe.eval()
        return dfe
    if model_path == "v4":
        dfe = DiffusionFeatureExtractor4(vae=vae, sd=sd)
        dfe.eval()
        return dfe
    if model_path == "v5":
        dfe = DiffusionFeatureExtractor5(vae=vae, sd=sd)
        dfe.eval()
        return dfe
    if model_path == "v6":
        dfe = DiffusionFeatureExtractor6(vae=vae, sd=sd)
        dfe.eval()
        return dfe
    if model_path == "v7":
        dfe = DiffusionFeatureExtractor7(vae=vae, sd=sd)
        dfe.eval()
        return dfe
    if model_path == "v8":
        dfe = DiffusionFeatureExtractor8(vae=vae, sd=sd)
        dfe.eval()
        return dfe
    if model_path == "v9":
        dfe = DiffusionFeatureExtractor9(vae=vae, sd=sd)
        dfe.eval()
        return dfe
    if model_path == "v10":
        dfe = DiffusionFeatureExtractor10(vae=vae, sd=sd)
        dfe.eval()
        return dfe
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model file not found: {model_path}")
    # if it ende with safetensors
    if model_path.endswith('.safetensors'):
        state_dict = load_file(model_path)
    else:
        state_dict = torch.load(model_path, weights_only=True)
        if 'model_state_dict' in state_dict:
            state_dict = state_dict['model_state_dict']

    if 'conv_in.weight' in state_dict:
        # determine num out channels
        out_channels = state_dict['conv_out.weight'].shape[0]
        dfe = DiffusionFeatureExtractor(out_channels=out_channels)
    else:
        dfe = DiffusionFeatureExtractor2()

    dfe.load_state_dict(state_dict)
    dfe.eval()
    return dfe