File size: 11,324 Bytes
c8c00f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 | from dataclasses import dataclass
import torch
from torch import Tensor, nn
import numpy as np
from flux.modules.layers import (DoubleStreamBlock, EmbedND, LastLayer,
MLPEmbedder, SingleStreamBlock,
timestep_embedding)
@dataclass
class FluxParams:
in_channels: int
out_channels: int
vec_in_dim: int
context_in_dim: int
hidden_size: int
mlp_ratio: float
num_heads: int
depth: int
depth_single_blocks: int
axes_dim: list[int]
theta: int
qkv_bias: bool
guidance_embed: bool
class Flux(nn.Module):
"""
Transformer model for flow matching on sequences.
"""
def __init__(self, params: FluxParams):
super().__init__()
self.params = params
self.in_channels = params.in_channels
self.out_channels = params.out_channels
if params.hidden_size % params.num_heads != 0:
raise ValueError(
f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
)
pe_dim = params.hidden_size // params.num_heads
if sum(params.axes_dim) != pe_dim:
raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
self.hidden_size = params.hidden_size
self.num_heads = params.num_heads
self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
self.guidance_in = (
MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity()
)
self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)
self.double_blocks = nn.ModuleList(
[
DoubleStreamBlock(
self.hidden_size,
self.num_heads,
mlp_ratio=params.mlp_ratio,
qkv_bias=params.qkv_bias,
)
for _ in range(params.depth)
]
)
self.single_blocks = nn.ModuleList(
[
SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio)
for _ in range(params.depth_single_blocks)
]
)
self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
self._sequential_offload = False
def enable_sequential_cpu_offload(self):
self._sequential_offload = True
def forward(
self,
img: Tensor,
img_ids: Tensor,
txt: Tensor,
txt_ids: Tensor,
timesteps: Tensor,
y: Tensor,
guidance: Tensor | None = None,
info = None,
ref_img: Tensor | None = None, # ← NEW
ref_img_ids: Tensor | None = None, # ← NEW
) -> Tensor:
if img.ndim != 3 or txt.ndim != 3:
raise ValueError("Input img and txt tensors must have 3 dimensions.")
# Ensure inputs match the model's dtype (NF4 can silently upcast to float32)
target_dtype = self.img_in.weight.dtype
img = img.to(target_dtype)
txt = txt.to(target_dtype)
if y.dtype != target_dtype:
y = y.to(target_dtype)
# running on sequences img
img = self.img_in(img)
original_img_seq_len = img.shape[1] # ← REMEMBER: how many original img tokens
vec = self.time_in(timestep_embedding(timesteps, 256))
if self.params.guidance_embed:
if guidance is None:
raise ValueError("Didn't get guidance strength for guidance distilled model.")
vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
vec = vec + self.vector_in(y)
txt = self.txt_in(txt)
# ═══════════════════════════════════════════════════════
# NEW: Concatenate reference tokens into image stream
# ═══════════════════════════════════════════════════════
if ref_img is not None and ref_img_ids is not None:
ref = self.img_in(ref_img) # project reference patches same way
img = torch.cat([img, ref], dim=1)
img_ids = torch.cat([img_ids, ref_img_ids], dim=1)
if ref_img is not None:
print(f"[Flux.forward] Attending to {ref_img.shape[1]} ref tokens + {original_img_seq_len} img tokens")
# ═══════════════════════════════════════════════════════
ids = torch.cat((txt_ids, img_ids), dim=1)
pe = self.pe_embedder(ids)
inject_pe = pe.clone()
if not info['inverse']:
# Defensive clamp: GSAM indices may exceed seq_len for non-16-divisible images
seq_len = pe.shape[2]
# Initialize accumulated lists for tracking all processed IDs
accumulated_target_ids = []
accumulated_ref_ids = []
for artifact_data in info['artifact_data']:
if artifact_data['artifact_type'] == 'addition' and info['addition']:
ref_ids = artifact_data['reference_patch_indices'].copy()
target_ids = artifact_data['target_patch_indices'].copy()
ref_ids = [max(0, min(int(i), seq_len - 1)) for i in ref_ids]
target_ids = [max(0, min(int(i), seq_len - 1)) for i in target_ids]
if len(target_ids) > 0 and len(ref_ids) > 0:
inject_pe[:,:,target_ids,:,:,:] = inject_pe[:,:,ref_ids,:,:,:]
# Accumulate IDs
if info['inject']:
accumulated_target_ids.extend(target_ids)
accumulated_ref_ids.extend(ref_ids)
elif artifact_data['artifact_type'] == 'removal' and info['removal']:
ref_ids = artifact_data['reference_patch_indices'].copy()
target_ids = artifact_data['target_patch_indices'].copy()
ref_ids = [max(0, min(int(i), seq_len - 1)) for i in ref_ids]
target_ids = [max(0, min(int(i), seq_len - 1)) for i in target_ids]
# ref_ids = get_closest_patch_inds(info['patch_h'], info['patch_w'], target_ids, ref_ids)
if len(target_ids) > 0 and len(ref_ids) > 0:
inject_pe[:,:,target_ids,:,:,:] = inject_pe[:,:,ref_ids,:,:,:]
# Accumulate IDs (after target_ids modification)
if info['inject']:
accumulated_target_ids.extend(target_ids)
accumulated_ref_ids.extend(ref_ids)
elif artifact_data['artifact_type'] == 'distortion' and info['distortion']:
ref_ids = artifact_data['reference_patch_indices'].copy()
target_ids = artifact_data['target_patch_indices'].copy()
ref_ids = [max(0, min(int(i), seq_len - 1)) for i in ref_ids]
target_ids = [max(0, min(int(i), seq_len - 1)) for i in target_ids]
if len(ref_ids) == 0:
# For distortion with no reference patches, shuffle target patches
ref_ids = target_ids.copy()
np.random.shuffle(ref_ids)
# Ensure target_ids and ref_ids are different for distortion
if len(target_ids) > 0 and len(ref_ids) > 0:
inject_pe[:,:,target_ids,:,:,:] = inject_pe[:,:,ref_ids,:,:,:]
# Accumulate IDs (after any ref_ids modification)
if info['inject']:
accumulated_target_ids.extend(target_ids)
accumulated_ref_ids.extend(ref_ids)
elif artifact_data['artifact_type'] == 'fusion' and info['fusion']:
ref_ids = artifact_data['reference_patch_indices'].copy()
target_ids = artifact_data['target_patch_indices'].copy()
ref_ids = [max(0, min(int(i), seq_len - 1)) for i in ref_ids]
target_ids = [max(0, min(int(i), seq_len - 1)) for i in target_ids]
# np.random.shuffle(ref_ids)
if len(target_ids) > 0 and len(ref_ids) > 0:
inject_pe[:,:,target_ids,:,:,:] = inject_pe[:,:,ref_ids,:,:,:]
# Accumulate IDs
if info['inject']:
accumulated_target_ids.extend(target_ids)
accumulated_ref_ids.extend(ref_ids)
info['patch_ids'] = accumulated_target_ids
info['patch_ref_ids'] = accumulated_ref_ids
info['timesteps'] = timesteps
if self._sequential_offload:
for block in self.double_blocks:
block = block.to(img.device)
img, txt = block(img=img, txt=txt, vec=vec, pe=inject_pe, info=info)
block = block.cpu()
torch.cuda.empty_cache()
else:
for block in self.double_blocks:
img, txt = block(img=img, txt=txt, vec=vec, pe=inject_pe, info=info)
cnt = 0
img = torch.cat((txt, img), 1)
info['type'] = 'single'
if self._sequential_offload:
for block in self.single_blocks:
block = block.to(img.device)
info['id'] = cnt
if cnt < 19:
img, info = block(img, vec=vec, pe=inject_pe, info=info)
else:
img, info = block(img, vec=vec, pe=pe, info=info)
block = block.cpu()
torch.cuda.empty_cache()
cnt += 1
else:
for block in self.single_blocks:
info['id'] = cnt
if cnt < 19:
img, info = block(img, vec=vec, pe=inject_pe, info=info)
else:
img, info = block(img, vec=vec, pe=pe, info=info)
cnt += 1
# ═══════════════════════════════════════════════════════
# MODIFIED: Extract only ORIGINAL img tokens
# Before: img = img[:, txt.shape[1] :, ...] (gets img + ref)
# After: img = img[:, txt.shape[1] : txt.shape[1] + original_img_seq_len, ...]
# ═══════════════════════════════════════════════════════
img = img[:, txt.shape[1] : txt.shape[1] + original_img_seq_len, ...]
img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)
return img, info |