Update autoregressive/models/generate.py

autoregressive/models/generate.py

@@ -90,9 +90,6 @@ def logits_to_probs(logits, temperature: float = 1.0, top_p: float=1.0, top_k: i
 
 def prefill(model, cond_idx: torch.Tensor, input_pos: torch.Tensor, cfg_scale: float, condition:torch.Tensor, **sampling_kwargs):
     if cfg_scale > 1.0:
-        print(cond_idx)
-        print(input_pos)
-        print(condition)
         logits, _ = model(None, cond_idx, input_pos, condition=condition)
         logits_combined = logits
         cond_logits, uncond_logits = torch.split(logits_combined, len(logits_combined) // 2, dim=0)
@@ -142,9 +139,11 @@ def decode_n_tokens(
 
 @torch.no_grad()
 def generate(model, cond, max_new_tokens, emb_masks=None, cfg_scale=1.0, cfg_interval=-1, condition=None, condition_null=None, condition_token_nums=0, **sampling_kwargs):
+    print(condition)
     if condition is not None:
         condition = model.adapter(condition)
         condition = model.adapter_mlp(condition)
+    print(condition)
     if model.model_type == 'c2i':
         if cfg_scale > 1.0:
             cond_null = torch.ones_like(cond) * model.num_classes