File size: 10,020 Bytes

99cb0de

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