pipeline.py

import safetensors.torch as st
import torch
from diffusers import StableDiffusionXLPipeline
from transformers import T5TokenizerFast, T5EncoderModel

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm.auto import tqdm

# ─────────────────────────────────────────────────────────────
# ░ Two-Stream Shunt Adapter
# ─────────────────────────────────────────────────────────────
class TwoStreamShuntAdapter(nn.Module):
    """
    Cross-attentive adapter that aligns T5 and CLIP token streams.

    Returns:
        anchor     : (B, Lc, clip_dim)
        delta      : (B, Lc, clip_dim)
        log_sigma  : (B, Lc, clip_dim)  – log σ, always finite
        attn_t2c   : (B, heads, Lt, Lc)
        attn_c2t   : (B, heads, Lc, Lt)
        tau        : (heads, 1, 1)       – per-head threshold param
        g_pred     : (B, 1)              – guidance-scale prediction
        gate       : (B, Lc, 1)          – per-token gate ∈ (0,1)
    """

    def __init__(
        self,
        t5_dim: int = 512,
        clip_dim: int = 768,
        bottleneck: int = 256,
        heads: int = 8,
        tau_init: float = 0.1,
        max_guidance: float = 10.0,
    ):
        super().__init__()
        print("TwoStreamShuntAdapter init")
        self.heads = heads
        self.bneck = bottleneck
        self.max_guidance = max_guidance

        # projections
        self.proj_t5   = nn.Linear(t5_dim,   bottleneck)
        self.proj_clip = nn.Linear(clip_dim, bottleneck)

        # cross-attention
        self.cross_t2c = nn.MultiheadAttention(
            bottleneck, heads, batch_first=True, dropout=0.1
        )
        self.cross_c2t = nn.MultiheadAttention(
            bottleneck, heads, batch_first=True, dropout=0.1
        )

        # head-wise τ
        self.tau = nn.Parameter(torch.full((heads, 1, 1), tau_init))

        # convolutional pocket residual (depth-wise)
        self.res1 = nn.Conv1d(
            bottleneck, bottleneck, 3, padding=1, groups=bottleneck
        )
        self.res2 = nn.Conv1d(
            bottleneck, bottleneck, 3, padding=1, groups=bottleneck
        )
        self.norm_res = nn.LayerNorm(bottleneck)

        # fusion + projections
        self.fuse = nn.Linear(2 * bottleneck, bottleneck)

        self.anchor_proj = nn.Sequential(
            nn.Linear(bottleneck, bottleneck), nn.GELU(),
            nn.Linear(bottleneck, clip_dim)
        )
        self.delta_proj = nn.Sequential(
            nn.Linear(bottleneck, bottleneck), nn.GELU(),
            nn.Linear(bottleneck, clip_dim)
        )
        self.logsig_proj = nn.Sequential(
            nn.Linear(bottleneck, bottleneck), nn.GELU(),
            nn.Linear(bottleneck, clip_dim)
        )
        self.gate_proj = nn.Sequential(
            nn.Linear(bottleneck, bottleneck), nn.GELU(),
            nn.Linear(bottleneck, 1), nn.Sigmoid()
        )
        self.guidance_proj = nn.Sequential(
            nn.LayerNorm(bottleneck), nn.Linear(bottleneck, 1), nn.Sigmoid()
        )

    def load_state_dict(self, args, **kwargs):
      # remove _orig_mod from state dict before applying.
      state_dict = {k.replace("_orig_mod.", ""): v for k, v in args.items()}
      super().load_state_dict(state_dict, **kwargs)

    def forward(self, t5_seq: torch.Tensor, clip_seq: torch.Tensor):
        print("📣  SHUNT FORWARD CALLED")

        B, Lt, _ = t5_seq.size()
        _, Lc, _ = clip_seq.size()

        # 1) project into bottleneck
        t5_b   = self.proj_t5(t5_seq)      # (B, Lt, b)
        clip_b = self.proj_clip(clip_seq)  # (B, Lc, b)

        # 2) cross-attention
        t2c, attn_t2c = self.cross_t2c(
            t5_b, clip_b, clip_b, need_weights=True, average_attn_weights=False
        )
        c2t, attn_c2t = self.cross_c2t(
            clip_b, t5_b, t5_b, need_weights=True, average_attn_weights=False
        )

        # 3) convolutional pocket on T5→CLIP
        x = t2c.transpose(1, 2)                    # (B, b, Lt)
        x = F.gelu(self.res1(x))
        x = F.gelu(self.res2(x)).transpose(1, 2)   # (B, Lt, b)
        pocket = self.norm_res(t2c + x)            # (B, Lt, b)

        # 4) fuse pocket avg with C2T
        pocket_mean = pocket.mean(1, keepdim=True).expand(-1, Lc, -1)
        h = F.gelu(self.fuse(torch.cat([pocket_mean, c2t], -1)))  # (B, Lc, b)

        # 5) outputs
        anchor     = self.anchor_proj(h)                       # (B,Lc,768)
        delta_mean = self.delta_proj(h)                        # (B,Lc,768)
        log_sigma  = self.logsig_proj(h)                       # (B,Lc,768)
        gate       = self.gate_proj(h)                         # (B,Lc,1)
        delta      = delta_mean * gate                         # (B,Lc,768)

        g_tok  = self.guidance_proj(h).squeeze(-1)             # (B,Lc)
        g_pred = g_tok.mean(1, keepdim=True) * self.max_guidance
        
        #print(anchor, delta, log_sigma, attn_t2c, attn_c2t, self.tau, g_pred, gate)

        return anchor, delta, log_sigma, attn_t2c, attn_c2t, self.tau, g_pred, gate

# --- 1. load pipeline -------------------------------------------------
pipe = StableDiffusionXLPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0",
        torch_dtype=torch.float16).to("cuda")

# --- 2. load tiny-T5 & shunt (fp32) -----------------------------------
t5_tok = T5TokenizerFast.from_pretrained("t5-small")
t5_mod = T5EncoderModel.from_pretrained("t5-small").eval().to("cuda")
shunt  = TwoStreamShuntAdapter().float().eval().to("cuda")
shunt.load_state_dict( st.load_file("/content/drive/MyDrive/t5-clip-l-shunts/vitl14_t5small_shunt_vanilla_final.safetensors") )

# --- 3. wrap encode_prompt once ---------------------------------------
orig_encode = pipe.encode_prompt

config = {
    "strength": 1.0,
    "gate_gamma": 1.0,
    "tau_scale": 1.0,
    "guidance_gain": 1.0,
    "guidance_bias": 0.0
}


gen = torch.Generator(device="cuda").manual_seed(420)



strength = 0

# the working version that can't be omitted, 
def stable_encode_prompt_shunted(self, *args, **kw):
    pe, ne, pool, npool = orig_encode(*args, **kw)   # regular call

    # 👉 split: first 768 dims are CLIP-L, rest 1280 are CLIP-G
    clipL, clipG = pe[..., :768], pe[..., 768:]

    # build T5 batch (handles CFG dup automatically because
    # encode_prompt already concatenated negative & positive if needed)
    bsz = clipL.shape[0]
    texts = ["tmp"] * bsz        # dummy, we only care about hidden states
    t5_ids  = t5_tok(texts, return_tensors="pt").input_ids.to("cuda")
    t5_seq  = t5_mod(t5_ids).last_hidden_state        # (B,L,512)

    # run adapter in fp32
    delta = shunt(t5_seq.float(), clipL.float())[1]   # second output is Δ
    delta = delta * strength                              # << your strength knob
    clipL_shift = (clipL.float() + delta).to(clipL.dtype)

    pe_shifted = torch.cat([clipL_shift, clipG], dim=-1)
    return pe_shifted, ne, pool, npool
#-----------------------------------------------------------------------------------------

def encode_prompt_shunted(self, *a, **k):
    # 1) run the normal encoder with “style” & “context” already split
    pe, ne, pool, npool = orig_encode(*a, **k)          # (B,77,2048)

    # 2) split CLIP-L / CLIP-G
    clipL, clipG = pe[..., :768], pe[..., 768:]

    # 3) build T5 on the *context* text (it’s in k['prompt_2'])
    t5_ids = t5_tok([k.get("prompt_2")], return_tensors="pt").input_ids.to(pe.device)
    t5_seq = t5_mod(t5_ids).last_hidden_state.float()

    # 4) shunt → Δ  (FP32 → back-cast)
    Δ = shunt(t5_seq, clipL.float())[1].to(clipL.dtype)
    clipL_shift = clipL + Δ * strength

    # 5) concatenate back
    pe_shift = torch.cat([clipL_shift, clipG], dim=-1)
    return pe_shift, ne, pool, npool

pipe.encode_prompt = encode_prompt_shunted.__get__(pipe, type(pipe))




PROMPT = "a naturally lit and beautiful room with a photorealistic depiction of a woman"
PROMPT_2 = "a realistic depiction of a woman sitting on a chair at a coffee shop sipping coffee, the environment is beautiful"
NEG = "blurry, distorted, monochrome, greyscale, watermark"
STEPS = 50
base_strength = 0.5
base_cfg = 7.5


for i in range(0, 4):
  strength = base_strength + (i * 0.25)
  cfg = base_cfg - (i * 0.25)
  img = pipe(
      PROMPT,
      prompt_2=PROMPT_2,
      negative_prompt=NEG, 
      num_inference_steps=STEPS,
      cfg_scale=cfg,
      generator=torch.Generator(device="cuda").manual_seed(420)
      ).images[0]
  img.save(f"woman_cfg_{int(cfg*100)}_{int(strength*100)}.png")

# --- 4. generate -------------------------------------------------------
#img = pipe(
#    PROMPT,
#    negative_prompt=NEG, 
#    num_inference_steps=STEPS,
#    generator=torch.Generator(device="cuda").manual_seed(420)
#    ).images[0]
#img.save("majestic_baseline.png")#
#

#strength = 0.25
## --- 4. generate -------------------------------------------------------
#img = pipe(
#    PROMPT,
#    negative_prompt=NEG, 
#    num_inference_steps=STEPS,
#    generator=torch.Generator(device="cuda").manual_seed(420)
#    ).images[0]
#img.save("majestic_02.png")#

#strength = 0.5
## --- 4. generate -------------------------------------------------------
#img = pipe(
#    PROMPT,
#    negative_prompt=NEG, 
#    num_inference_steps=STEPS,
#    generator=torch.Generator(device="cuda").manual_seed(420)
#    ).images[0]
#img.save("majestic_05.png")#

#strength = 0.75
## --- 4. generate -------------------------------------------------------
#img = pipe(
#    PROMPT,
#    negative_prompt=NEG, 
#    num_inference_steps=STEPS,
#    generator=torch.Generator(device="cuda").manual_seed(420)
#    ).images[0]
#img.save("majestic_075.png")