Upload app.py

app.py

@@ -75,16 +75,30 @@ def build_staircase_attention_mask(x, block_size, pad_id):
 
 
 def clone_past_key_values(pkv):
-    """Shallow clone of prefix KV-cache to avoid expensive copy.deepcopy."""
+    """Clone KV-cache. Fast path for tuples and Cache objects; falls back to deepcopy."""
     if pkv is None:
         return None
+    # Fast path: legacy tuple format
     if isinstance(pkv, tuple):
         return tuple(
             (k.clone() if k is not None else None, v.clone() if v is not None else None)
             for k, v in pkv
         )
-    return pkv
-
+    # Fast path: transformers Cache objects (DynamicCache, etc.)
+    if hasattr(pkv, 'key_cache') and hasattr(pkv, 'value_cache'):
+        try:
+            new_cache = pkv.__class__()
+            new_cache.key_cache = [k.clone() for k in pkv.key_cache]
+            new_cache.value_cache = [v.clone() for v in pkv.value_cache]
+            for attr in ('_seen_tokens', 'seen_tokens'):
+                if hasattr(pkv, attr):
+                    setattr(new_cache, attr, getattr(pkv, attr))
+            return new_cache
+        except Exception:
+            pass
+    # Fallback
+    import copy
+    return copy.deepcopy(pkv)
 
 def diffusion_step_block(logits, x_block, mask_block, num_transfer, temperature, remasking):
     """Vectorized diffusion step — no per-sample Python loops."""
@@ -172,18 +186,20 @@ def generate(model, tokenizer, prompt, steps=128, max_new_tokens=128, block_size
         # Clone prefix caches once per block instead of every step
         cond_past_clone = clone_past_key_values(cond_past)
         uncond_past_clone = clone_past_key_values(uncond_past) if uncond_past is not None else None
+        
         for t in range(eff_steps):
             x_blk = x[:, T_prefix:T_total]
             m_blk = x_blk == mask_id
             cond_logits = model(
                 x_blk, attention_mask=attn_blk, position_ids=pos_blk,
-                past_key_values=cond_past_clone, use_cache=False
+                past_key_values=clone_past_key_values(cond_past), use_cache=False
             ).logits
             logits = cond_logits
+            
             if cfg_scale > 0:
                 un_logits = model(
                     x_blk, attention_mask=attn_blk, position_ids=pos_blk,
-                    past_key_values=uncond_past_clone, use_cache=False
+                    past_key_values=clone_past_key_values(uncond_past), use_cache=False
                 ).logits
                 logits = un_logits + (cfg_scale + 1.0) * (cond_logits - un_logits)
             x_blk_new = diffusion_step_block(
@@ -267,13 +283,13 @@ def generate_stream(model, tokenizer, prompt, steps=128, max_new_tokens=128, blo
             m_blk = x_blk == mask_id
             cond_logits = model(
                 x_blk, attention_mask=attn_blk, position_ids=pos_blk,
-                past_key_values=cond_past_clone, use_cache=False
+                past_key_values=clone_past_key_values(cond_past), use_cache=False
             ).logits
             logits = cond_logits
             if cfg_scale > 0:
                 un_logits = model(
                     x_blk, attention_mask=attn_blk, position_ids=pos_blk,
-                    past_key_values=uncond_past_clone, use_cache=False
+                    past_key_values=clone_past_key_values(uncond_past), use_cache=False
                 ).logits
                 logits = un_logits + (cfg_scale + 1.0) * (cond_logits - un_logits)
             x_blk_new = diffusion_step_block(