| The 155000 step version has about 158,100,000 prompt samples weight trained using the | |
| AbstractPhil/T5-Small-Human-Attentive-Try2-Pass3 | |
| This T5-small model is fried to echo and interpolate math in complex intended ways. I haven't given it the full robust check yet, but it's definitely pretty fed. | |
| This adapter here is trained using T5 inputs with this code below. | |
| This isn't a bad first test. I will be improving the adapter with common lora techniques, including more techniques from training LLM-style loras, and including additional loss methodologies while simultaneously including more advanced and carefully curated response formulas to the way the adapter responded to training and the extrapolative math from the CLIP_L adapted response. | |
| Given time I'm certain this will work; whether it be creating a layered lora structure to interpolate differences layer by layer within the clip_l, or perhaps in a much more direct neuron interpolation. Time will tell and I'm definitely enjoying this sort of thing. | |
| Errors to address in the next; | |
| * There is a clamping index error that tends to rear it's head that I haven't had time to track down. It'll cause solid black images from the velocity sigmas being too heavy. | |
| * Occasionally the entire structure of a generation collapses, which means the sigmas aren't lined up correctly - creating malformed sigma responses. | |
| * Occasionally the substructure interprets the request incorrectly; this is due to the tokenization being inaccurately attuned for some spaces than others and this next version will have node weighting for specific attention head sectors to account for it. | |
| There's many challenges ahead to reach the interpolation endpoint but it's definitely an adaptive journey. | |
| This is stage 1 of multiple stages to make the recreatable pragmatic outcomes needed in order to build the proofs required to recreate the Beatrix interpolation model - into useful utiliizations outside of diffusion. | |
| This process adapts multiple similar methods as what I used to create the Beatrix model, but it's not 1:1 by any stretch of the measure. | |
| I will be slowly releasing parts of Beatrix in training diagrams and stage the methodologies about how she works, so the interested experts will be capable of rationalizing why this model does what it does. | |
| Because I really don't know why Beatrix works the way she does, and I'm not going to just release something like that until I understand WHY it skips and hops past entropy. | |
| 77 tokens - not 64, there's no need to upscale the most recent 77tok version; it's built to the same plane as CLIP_L now. | |
| ``` | |
| def main(): | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # HF Hub settings | |
| hf_repo_id = "AbstractPhil/t5-to-vit-l-14-velocity-adapter-v3-100m-77tok" | |
| push_every_n_steps = 5000 | |
| # Tokenizers & frozen models | |
| t5_tok = T5TokenizerFast.from_pretrained("t5-small") | |
| t5_mod = T5EncoderModel.from_pretrained( | |
| "AbstractPhil/T5-Small-Human-Attentive-Try2-Pass3" | |
| ).to(device).eval() | |
| clip_tok = CLIPTokenizerFast.from_pretrained("openai/clip-vit-large-patch14") | |
| clip_mod = CLIPTextModel.from_pretrained( | |
| "openai/clip-vit-large-patch14" | |
| ).to(device).eval() | |
| # Adapter & optimizer | |
| adapter = RobustVelocityAdapter(out_tokens=77).to(device) | |
| optimizer = optim.AdamW(adapter.parameters(), lr=5e-4) | |
| # Compile models for speed | |
| t5_mod = torch.compile(t5_mod) | |
| clip_mod = torch.compile(clip_mod) | |
| adapter = torch.compile(adapter) | |
| scaler = GradScaler() # for mixed precision | |
| # Data | |
| dataset = ParsedMultiCharDataset("AbstractPhil/human-templated-captions-1b", | |
| num_files=12) | |
| loader = DataLoader(dataset, | |
| batch_size=None, | |
| num_workers=4, | |
| pin_memory=True) | |
| iterator = iter(loader) | |
| batch_size = 256 | |
| accum_steps = 4 # effective BS = 256 * 4 = 1024 | |
| max_steps = math.ceil(dataset.total_rows / batch_size) | |
| pbar = tqdm(total=max_steps, desc="Adapter training") | |
| for step in range(1, max_steps+1): | |
| # zero grads on actual step | |
| if (step-1) % accum_steps == 0: | |
| optimizer.zero_grad() | |
| # 1) Collect batch | |
| texts = [] | |
| for _ in range(batch_size): | |
| try: | |
| _, txt = next(iterator) | |
| except StopIteration: | |
| iterator = iter(loader) | |
| _, txt = next(iterator) | |
| texts.append(txt) | |
| # 2) Tokenize | |
| t5_inputs = t5_tok(texts, | |
| padding=True, | |
| truncation=True, | |
| max_length=77, | |
| return_tensors="pt").to(device) | |
| clip_inputs = clip_tok(texts, | |
| padding="max_length", | |
| truncation=True, | |
| max_length=77, | |
| return_tensors="pt").to(device) | |
| # 3) Forward + loss in mixed precision | |
| with autocast(): | |
| t5_seq = t5_mod(**t5_inputs).last_hidden_state # [B,77,512] | |
| clip_seq = clip_mod(**clip_inputs).last_hidden_state # [B,77,768] | |
| anchor_pred, delta_pred, sigma_pred = adapter(t5_seq) | |
| delta_target = clip_seq - anchor_pred | |
| loss_delta = hetero_loss(delta_pred, delta_target, sigma_pred) | |
| # cosine anchor alignment | |
| cos_sim = nn.functional.cosine_similarity( | |
| anchor_pred.reshape(-1,768), | |
| clip_seq.reshape(-1,768), | |
| dim=-1 | |
| ).mean() | |
| loss_anchor = (1 - cos_sim) * 0.1 | |
| loss = loss_delta + loss_anchor | |
| loss = loss / accum_steps # scale for accumulation | |
| # 4) Backward + optimizer step | |
| scaler.scale(loss).backward() | |
| if step % accum_steps == 0: | |
| scaler.unscale_(optimizer) | |
| torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0) | |
| scaler.step(optimizer) | |
| scaler.update() | |
| pbar.update(1) | |
| pbar.set_postfix(loss=(loss.item() * accum_steps)) | |
| # 5) Save & push every N steps | |
| if step % push_every_n_steps == 0: | |
| ckpt = f"/content/drive/MyDrive/t5-adapter/t5-to-vit-l-14-velocity-adapter-v3-100m-77tok_step_{step}.safetensors" | |
| save_file(adapter.state_dict(), ckpt) | |
| #upload_file(ckpt, ckpt, repo_id=hf_repo_id) | |
| pbar.close() | |
| ``` | |
| ### You'll need to snip out the __orig layer extensions that got snapped into it when I saved. | |
| Still not quite sure how to fix that without just editing before saving, but I think it's causing some sort of additional effects that I'm unaware of. | |
| I don't want to save as pt because they are considered unsafe and I don't want this to be considered unsafe for use. | |
| You can inference the test version using stable-diffusion-15 as an example test. | |
| The CLIP_L responses fall apart when too many nodes hit those guidance bells, but it's definitely a powerful first test using divergent systems. | |
| Should just run clean on colab using a l4. | |
| ``` | |
| # Optimized inference_adapter.py | |
| import torch | |
| import math | |
| from PIL import Image | |
| from torchvision.transforms import ToPILImage | |
| from safetensors.torch import load_file as load_safetensors | |
| from transformers import ( | |
| T5TokenizerFast, T5EncoderModel, | |
| CLIPTokenizerFast, CLIPTextModel | |
| ) | |
| from diffusers import ( | |
| AutoencoderKL, | |
| UNet2DConditionModel, | |
| EulerAncestralDiscreteScheduler | |
| ) | |
| from typing import Optional | |
| # βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ | |
| # 1) GLOBAL SETUP: load once, cast, eval, move | |
| # βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ | |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| DTYPE = torch.float16 # use fp16 for everything on GPU | |
| # 1a) CLIP text encoder (cond + uncond) | |
| clip_tok = CLIPTokenizerFast.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", subfolder="tokenizer" | |
| ) | |
| clip_mod = CLIPTextModel.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", subfolder="text_encoder", | |
| torch_dtype=DTYPE | |
| ).to(DEVICE).eval() | |
| # 1b) T5 encoder | |
| t5_tok = T5TokenizerFast.from_pretrained("t5-small") | |
| t5_mod = T5EncoderModel.from_pretrained( | |
| "AbstractPhil/T5-Small-Human-Attentive-Try2-Pass3", | |
| torch_dtype=DTYPE | |
| ).to(DEVICE).eval() | |
| # 1c) Velocity Adapter local directory | |
| local_adapter_directory = "roba_adapter_step_19500.safetensors" # opens the state below. | |
| # 1c) Adapter | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class RobustVelocityAdapter(nn.Module): | |
| """ | |
| Fixed version: manual multi-head cross-attention emits [B, heads, Q, K] scores | |
| so that _add_rel_pos_bias can unpack them correctly. | |
| """ | |
| def __init__( | |
| self, | |
| t5_dim: int = 512, | |
| clip_dim: int = 768, | |
| hidden_dim: int = 1024, | |
| out_tokens: int = 77, # now aligned with your T5 finetune | |
| self_attn_layers: int = 2, | |
| cross_heads: int = 8, | |
| max_rel_pos: int = 128, | |
| ): | |
| super().__init__() | |
| self.out_tokens = out_tokens | |
| self.cross_heads = cross_heads | |
| self.head_dim = t5_dim // cross_heads | |
| self.max_rel_pos = max_rel_pos | |
| # 1) Self-attention stack | |
| self.self_attn = nn.ModuleList() | |
| self.self_norm = nn.ModuleList() | |
| for _ in range(self_attn_layers): | |
| self.self_attn.append(nn.MultiheadAttention(t5_dim, cross_heads, batch_first=True)) | |
| self.self_norm.append(nn.LayerNorm(t5_dim)) | |
| # 2) Residual blocks | |
| def resblock(): | |
| return nn.Sequential( | |
| nn.LayerNorm(t5_dim), | |
| nn.Linear(t5_dim, t5_dim), | |
| nn.GELU(), | |
| nn.Linear(t5_dim, t5_dim), | |
| ) | |
| self.res1 = resblock() | |
| self.res2 = resblock() | |
| # 3) Learned queries for cross-attn | |
| self.query_pos = nn.Parameter(torch.randn(out_tokens, t5_dim)) | |
| # 4) Projection heads | |
| self.anchor_proj = nn.Sequential( | |
| nn.Linear(t5_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, clip_dim) | |
| ) | |
| self.delta_proj = nn.Sequential( | |
| nn.Linear(t5_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, clip_dim) | |
| ) | |
| self.var_proj = nn.Sequential( | |
| nn.Linear(t5_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, clip_dim) | |
| ) | |
| self.gate_proj = nn.Sequential( | |
| nn.Linear(t5_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, clip_dim), nn.Sigmoid() | |
| ) | |
| # 5) Relative-position bias table | |
| self.rel_bias = nn.Parameter(torch.zeros(2*max_rel_pos-1, cross_heads)) | |
| # 6) Norm after cross-attn | |
| self.cross_norm = nn.LayerNorm(t5_dim) | |
| def _add_rel_pos_bias(self, attn_scores: torch.Tensor) -> torch.Tensor: | |
| """ | |
| attn_scores: [B, heads, Q, K] | |
| returns: attn_scores + bias where bias is [B, heads, Q, K] | |
| """ | |
| B, H, Q, K = attn_scores.shape | |
| device = attn_scores.device | |
| # 1) Query & key position indices | |
| idx_q = torch.arange(Q, device=device) # [Q] | |
| idx_k = torch.arange(K, device=device) # [K] | |
| # 2) Compute relative distances for every (q, k) pair | |
| # rel[i,j] = idx_q[i] - idx_k[j] | |
| rel = idx_q.unsqueeze(1) - idx_k.unsqueeze(0) # [Q, K] | |
| # 3) Clamp & shift into bias table range [0, 2*max_rel-2] | |
| max_rel = self.max_rel_pos | |
| rel = rel.clamp(-max_rel+1, max_rel-1) + (max_rel - 1) | |
| # 4) Lookup per-head biases | |
| # self.rel_bias has shape [2*max_rel-1, H] | |
| bias = self.rel_bias[rel] # [Q, K, H] | |
| bias = bias.permute(2, 0, 1) # [H, Q, K] | |
| # 5) Broadcast to [B, H, Q, K] and add | |
| bias = bias.unsqueeze(0).expand(B, -1, -1, -1) | |
| return attn_scores + bias | |
| def forward(self, t5_seq: torch.Tensor): | |
| """ | |
| t5_seq: [B, L, t5_dim] | |
| returns: | |
| anchor: [B, out_tokens, clip_dim] | |
| delta: [B, out_tokens, clip_dim] | |
| sigma: [B, out_tokens, clip_dim] | |
| """ | |
| x = t5_seq | |
| B, L, D = x.shape | |
| # 1) Self-attention + residual | |
| for attn, norm in zip(self.self_attn, self.self_norm): | |
| res, _ = attn(x, x, x) | |
| x = norm(x + res) | |
| # 2) Residual blocks | |
| x = x + self.res1(x) | |
| x = x + self.res2(x) | |
| # 3) Prepare queries & split heads | |
| queries = self.query_pos.unsqueeze(0).expand(B, -1, -1) # [B, Q, D] | |
| # reshape into heads | |
| q = queries.view(B, self.out_tokens, self.cross_heads, self.head_dim).permute(0,2,1,3) | |
| k = x.view(B, L, self.cross_heads, self.head_dim).permute(0,2,1,3) | |
| v = k | |
| # 4) Scaled dot-product to get [B, heads, Q, K] | |
| scores = (q @ k.transpose(-2,-1)) / math.sqrt(self.head_dim) | |
| scores = self._add_rel_pos_bias(scores) | |
| probs = F.softmax(scores, dim=-1) # [B, H, Q, K] | |
| # 5) Attend & merge heads β [B, Q, D] | |
| ctx = probs @ v # [B, H, Q, head_dim] | |
| ctx = ctx.permute(0,2,1,3).reshape(B, self.out_tokens, D) | |
| ctx = self.cross_norm(ctx) | |
| # 6) Project to anchor, delta_mean, delta_logvar, gate | |
| anchor = self.anchor_proj(ctx) | |
| delta_mean = self.delta_proj(ctx) | |
| delta_logvar = self.var_proj(ctx) | |
| gate = self.gate_proj(ctx) | |
| # 7) Compute sigma & gated delta | |
| sigma = torch.exp(0.5 * delta_logvar) | |
| delta = delta_mean * gate | |
| return anchor, delta, sigma | |
| import torch | |
| import torch.nn.functional as F | |
| from PIL import Image | |
| from torchvision.transforms import ToPILImage | |
| from safetensors.torch import load_file as load_safetensors | |
| from transformers import ( | |
| CLIPTokenizer, CLIPTextModel, | |
| T5TokenizerFast, T5EncoderModel | |
| ) | |
| from diffusers import ( | |
| AutoencoderKL, | |
| UNet2DConditionModel, | |
| EulerAncestralDiscreteScheduler | |
| ) | |
| # 1) GLOBAL SETUP | |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| DTYPE = torch.float32 | |
| # 1a) CLIP tokenizer & text encoder | |
| clip_tok = CLIPTokenizer.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", subfolder="tokenizer" | |
| ) | |
| clip_mod = CLIPTextModel.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", subfolder="text_encoder", | |
| torch_dtype=DTYPE | |
| ).to(DEVICE).eval() | |
| # 1b) U-Net, VAE, Scheduler | |
| unet = UNet2DConditionModel.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", subfolder="unet", | |
| torch_dtype=DTYPE | |
| ).to(DEVICE).eval() | |
| vae = AutoencoderKL.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", subfolder="vae", | |
| torch_dtype=DTYPE | |
| ).to(DEVICE).eval() | |
| scheduler = EulerAncestralDiscreteScheduler.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", subfolder="scheduler" | |
| ) | |
| # 1c) T5 | |
| t5_tok = T5TokenizerFast.from_pretrained("t5-small") | |
| t5_mod = T5EncoderModel.from_pretrained( | |
| "AbstractPhil/T5-Small-Human-Attentive-Try2-Pass3", | |
| torch_dtype=DTYPE | |
| ).to(DEVICE).eval() | |
| # 1d) velocity prediction adapter | |
| adapter = RobustVelocityAdapter(out_tokens=77).to(DEVICE).eval() | |
| state = load_safetensors(local_adapter_directory, device="cpu") | |
| clean = {k.replace("_orig_mod.", ""): v for k, v in state.items()} | |
| adapter.load_state_dict(clean, strict=False) | |
| adapter.to(DEVICE).eval() | |
| # 2) GENERATION FUNCTION | |
| @torch.no_grad() | |
| def generate_image_with_adapter( | |
| prompt: str, | |
| seed: int = 42, | |
| steps: int = 50, | |
| adapter_scale: float = 0.5, | |
| guidance_scale: float = 7.5, | |
| height: int = 512, | |
| width: int = 512, | |
| ): | |
| gen = torch.Generator(device=DEVICE).manual_seed(seed) | |
| # 2.1) CLIP embeddings | |
| clip_in = clip_tok([prompt], | |
| max_length=clip_tok.model_max_length, | |
| padding="max_length", truncation=True, | |
| return_tensors="pt").to(DEVICE) | |
| clip_cond = clip_mod(**clip_in).last_hidden_state # [1,77,768] | |
| empty_in = clip_tok([""], | |
| max_length=clip_tok.model_max_length, | |
| padding="max_length", truncation=True, | |
| return_tensors="pt").to(DEVICE) | |
| clip_uncond= clip_mod(**empty_in).last_hidden_state # [1,77,768] | |
| # 2.2) T5 β adapter β anchor, delta, sigma (77 tokens) | |
| t5_in = t5_tok(prompt, | |
| max_length=77, padding="max_length", | |
| truncation=True, return_tensors="pt").to(DEVICE) | |
| t5_seq = t5_mod(**t5_in).last_hidden_state # [1,77,512] | |
| anchor, delta, sigma = adapter(t5_seq) # each [1,77,768] | |
| # 2.3) Upsample to 77 tokens | |
| T_clip = clip_cond.shape[1] # 77 | |
| def up(x): | |
| return F.interpolate( | |
| x.permute(0,2,1), | |
| size=T_clip, mode="linear", align_corners=False | |
| ).permute(0,2,1) | |
| anchor = up(anchor) | |
| delta = up(delta) | |
| sigma = up(sigma) | |
| # 2.4) Ο-based noise scaling | |
| raw_ns = sigma.mean().clamp(0.1, 2.0).item() | |
| noise_scale = 1.0 + adapter_scale * (raw_ns - 1.0) | |
| # 2.5) Initialize latents | |
| latents = torch.randn( | |
| (1, unet.config.in_channels, height//8, width//8), | |
| generator=gen, device=DEVICE, dtype=DTYPE | |
| ) * scheduler.init_noise_sigma * noise_scale | |
| scheduler.set_timesteps(steps, device=DEVICE) | |
| # 2.6) Denoising with adapter guidance | |
| for i, t in enumerate(scheduler.timesteps): | |
| alpha = i / (len(scheduler.timesteps)-1) | |
| aw = adapter_scale * alpha | |
| cw = 1.0 - aw | |
| # blend anchors | |
| blended = clip_cond * cw + anchor * aw | |
| # per-token confidence | |
| eps = 1e-6 | |
| conf = 1.0 / (sigma + eps) | |
| conf = conf / conf.amax(dim=(1,2), keepdim=True) | |
| # gated delta | |
| gated_delta = delta * aw * conf | |
| # final cond embedding | |
| cond_embed = blended + gated_delta # [1,77,768] | |
| # UNet forward | |
| lat_in = scheduler.scale_model_input(latents, t) | |
| lat_in = torch.cat([lat_in, lat_in], dim=0) | |
| embeds = torch.cat([clip_uncond, cond_embed], dim=0) | |
| noise = unet(lat_in, t, encoder_hidden_states=embeds).sample | |
| u, c = noise.chunk(2) | |
| guided = u + guidance_scale * (c - u) | |
| latents= scheduler.step(guided, t, latents, generator=gen).prev_sample | |
| # 2.7) Decode | |
| dec_lat = latents / vae.config.scaling_factor | |
| image_t = vae.decode(dec_lat).sample | |
| image_t = (image_t.clamp(-1,1) + 1) / 2 | |
| return ToPILImage()(image_t[0]) | |
| # 3) RUN EXAMPLE | |
| if __name__ == "__main__": | |
| out = generate_image_with_adapter( | |
| "silly dog wearing a batman costume, high resolution, studio lighting", | |
| seed=1234, steps=50, | |
| adapter_scale=0.5, guidance_scale=7.5 | |
| ) | |
| out.save("sd15_with_adapter.png") | |
| print("Saved sd15_with_adapter.png") | |