Upload app.py

app.py

@@ -3,6 +3,7 @@ import types
 import os
 import math
 import json
+import copy
 import torch
 import torch.nn.functional as F
 from flask import Flask, request, jsonify, Response
@@ -97,9 +98,9 @@ def clone_past_key_values(pkv):
         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."""
     B, L, _ = logits.shape
@@ -183,10 +184,6 @@ def generate(model, tokenizer, prompt, steps=128, max_new_tokens=128, block_size
         full_attn, full_pos = build_staircase_attention_mask(x, block_size, pad_id)
         attn_blk = full_attn[:, :, T_prefix:T_total, :]
         pos_blk = full_pos[:, T_prefix:T_total]
-        # 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
@@ -195,7 +192,6 @@ def generate(model, tokenizer, prompt, steps=128, max_new_tokens=128, block_size
                 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,
@@ -275,9 +271,6 @@ def generate_stream(model, tokenizer, prompt, steps=128, max_new_tokens=128, blo
         full_attn, full_pos = build_staircase_attention_mask(x, block_size, pad_id)
         attn_blk = full_attn[:, :, T_prefix:T_total, :]
         pos_blk = full_pos[:, T_prefix:T_total]
-        # Clone prefix caches once per block
-        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
@@ -351,9 +344,6 @@ def load_model():
             print(f"torch.compile skipped: {e}")
     else:
         print("Diffusion model loaded without torch.compile (custom FX code incompatible with Dynamo).")
-
-
-
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
         trust_remote_code=IS_DIFFUSION