add 4 step turbo lora

app.py

@@ -9,6 +9,9 @@ stay in sync automatically.
 
 Both checkpoints are loaded once at module scope (ZeroGPU packs and transfers
 them when a GPU is actually allocated) and shared by every run of the node.
+Model choices are 'base', 'sft', and 'turbo' — the 4-step distilled LoRA from
+akhaliq/MOSS-Image-SFT-4-steps wrapped on the SFT transformer and toggled per
+request (see turbo_lora.py).
 """
 
 import os
@@ -157,13 +160,21 @@ def _load_pipeline(model_id, shared_pipe=None):
     snapshot = snapshot_download(
         model_id, revision=revision, ignore_patterns=[f"{name}/*" for name in components]
     )
+    transformer = MOSSImageTransformer.from_pretrained(
+        snapshot, subfolder="transformer", torch_dtype=torch.bfloat16
+    )
+    if model_id == MODEL_IDS["sft"]:
+        # Wrap the SFT transformer with the 4-step turbo LoRA (fp32, unmerged)
+        # BEFORE the pipeline moves to CUDA, so ZeroGPU packs it with the rest.
+        # Per request it is toggled by alpha (0 = off) — see turbo_lora.py.
+        from turbo_lora import apply_turbo_lora
+
+        apply_turbo_lora(transformer)
     pipe = MOSSImagePipeline.from_pretrained(
         snapshot,
         custom_pipeline=str(Path(__file__).with_name("pipeline_moss_image")),
         trust_remote_code=True,
-        transformer=MOSSImageTransformer.from_pretrained(
-            snapshot, subfolder="transformer", torch_dtype=torch.bfloat16
-        ),
+        transformer=transformer,
         scheduler=MOSSImageScheduler.from_pretrained(snapshot, subfolder="scheduler"),
         local_files_only=True,
         torch_dtype={"default": torch.bfloat16, "vae": torch.float32},
@@ -234,8 +245,11 @@ def generate(
     if not (prompt or "").strip():
         raise gr.Error("Enter a prompt.")
     key = (model or "").strip().lower().replace("moss-image-", "")
+    turbo = key in ("turbo", "4-step", "4-steps", "sft-turbo")
+    if turbo:
+        key = "sft"
     if key not in PIPELINES:
-        raise gr.Error(f"Unknown model {model!r}. Choose 'base' or 'sft'.")
+        raise gr.Error(f"Unknown model {model!r}. Choose 'base', 'sft' or 'turbo'.")
     width, height, steps, num_images, max_sequence_length = (
         round(width), round(height), round(steps), round(num_images),
         round(max_sequence_length),
@@ -257,6 +271,16 @@ def generate(
         raise gr.Error("Seed must be -1 (random) or an integer from 0 to 2^53 - 1.")
 
     pipe = PIPELINES[key]
+    if key == "sft":
+        from turbo_lora import set_turbo
+
+        set_turbo(pipe, turbo)
+    if turbo:
+        # Distilled model: CFG 1 always, and the published step counts unless
+        # the user picked their own (50 is the untouched default).
+        cfg = 1.0
+        if steps == 50:
+            steps = 4 if max(width, height) <= 1024 else 6 if max(width, height) <= 1536 else 8
     depth = pipe.transformer.config.num_layers
     skip = _parse_layers(uncond_skip_layers, "Skip layers")
     slg = _parse_layers(slg_layers, "SLG layers")
@@ -293,8 +317,9 @@ def generate(
         image.save(path)
         results.append(_as_media(path))
     elapsed = time.monotonic() - start
+    display = "MOSS-Image-SFT-4-steps" if turbo else MODEL_IDS[key].rsplit("/", 1)[-1]
     status = (
-        f"{len(results)} image(s) · {MODEL_IDS[key].rsplit('/', 1)[-1]} "
+        f"{len(results)} image(s) · {display} "
         f"· seed {seed} · {width}×{height} · {steps} steps "
         f"· CFG {cfg:g} · shift {shift:g} · {elapsed:.1f}s"
     )

turbo_lora.py

@@ -0,0 +1,76 @@
+"""Apply and toggle the MOSS-Image-SFT-4-steps turbo LoRA.
+
+Vendored from https://huggingface.co/akhaliq/MOSS-Image-SFT-4-steps
+(apply_lora.py, Apache-2.0), with a per-request on/off toggle added: the fp32
+LoRA stays wrapped on the SFT transformer, and `set_turbo` switches its alpha
+between the trained value and zero — so SFT and turbo generations share one
+pipeline in GPU memory instead of loading a third transformer.
+
+Usage matches the model card: 4 steps at CFG 1.0 (6–8 steps above 1024²).
+"""
+
+import torch
+from torch import nn
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+REPO = "akhaliq/MOSS-Image-SFT-4-steps"
+LORA_TARGETS = ("q", "k", "v", "o", "ffn_up_gate", "ffn_down")
+RANK = 128
+
+
+class LoRALinear(nn.Module):
+    """fp32 LoRA kept unmerged on the (bf16) base: y = W x + s * up(down(x))."""
+
+    def __init__(self, base, rank, alpha):
+        super().__init__()
+        self.base = base
+        for p in self.base.parameters():
+            p.requires_grad_(False)
+        self.lora_down = nn.Linear(base.in_features, rank, bias=False)
+        self.lora_up = nn.Linear(rank, base.out_features, bias=False)
+        self.rank, self.alpha = rank, alpha
+
+    @property
+    def weight(self):
+        return self.base.weight
+
+    def forward(self, x):
+        out = self.base(x)
+        if self.alpha == 0:
+            return out
+        delta = self.lora_up(self.lora_down(x.to(self.lora_down.weight.dtype)))
+        return out + (delta * (self.alpha / self.rank)).to(out.dtype)
+
+
+def apply_turbo_lora(transformer, repo=REPO, strength=1.0, filename="lora.safetensors"):
+    """Wrap every block's proj.{q,k,v,o,ffn_up_gate,ffn_down} with the fp32 LoRA.
+
+    filename: canonical "lora.safetensors" (always the latest release) or a
+    frozen versioned name, e.g. "moss_image_sft_4step_v2.0.safetensors".
+    """
+    for layer in transformer.layers:
+        for name in LORA_TARGETS:
+            layer.proj[name] = LoRALinear(layer.proj[name], RANK, RANK * strength)
+    lora_sd = load_file(hf_hub_download(repo, filename))
+    for i, layer in enumerate(transformer.layers):
+        for name in LORA_TARGETS:
+            m = layer.proj[name]
+            key = f"diffusion_model.layers.{i}.proj.{name}"
+            m.lora_down.weight.data.copy_(lora_sd[f"{key}.lora_down.weight"])
+            m.lora_up.weight.data.copy_(lora_sd[f"{key}.lora_up.weight"])
+            m.trained_alpha = float(lora_sd[f"{key}.alpha"]) * strength
+            m.alpha = m.trained_alpha
+    return transformer
+
+
+def set_turbo(pipe, enabled):
+    """Toggle the wrapped LoRA without reloading weights (alpha 0 = off).
+
+    A no-op for pipelines whose transformer was never wrapped (e.g. Base).
+    """
+    for layer in pipe.transformer.layers:
+        for name in LORA_TARGETS:
+            module = layer.proj[name]
+            if isinstance(module, LoRALinear):
+                module.alpha = module.trained_alpha if enabled else 0.0