Upload DUO

model.py

@@ -530,8 +530,11 @@ class HFDIT(torch.nn.Module):
     else:
       return  bias_dropout_add_scale_fused_inference
 
-  def forward(self, x, sigma):
+  def forward(self, x, sigma, output_hidden_states=False):
+    all_hidden_states = []
     x = self.vocab_embed(x)
+    if output_hidden_states:
+      all_hidden_states.append(x)
     if self.causal:
       t_cond = None
     else:
@@ -541,8 +544,12 @@ class HFDIT(torch.nn.Module):
     with torch.cuda.amp.autocast(dtype=torch.bfloat16):
       for i in range(len(self.blocks)):
         x = self.blocks[i](x, rotary_cos_sin, c=t_cond)
+        if output_hidden_states:
+          all_hidden_states.append(x)
       x = self.output_layer(x, c=t_cond)
-    return x
+      if output_hidden_states:
+        all_hidden_states.append(x)
+    return x, all_hidden_states
 
 
 
@@ -585,10 +592,8 @@ class DUO(transformers.PreTrainedModel):
       else self.config.use_return_dict
     
     logits, all_hidden_states = self.backbone(
-      indices=input_ids,
+      x=input_ids,
       sigma=timesteps,
-      sample_mode=sample_mode,
-      store_kv=store_kv,
       output_hidden_states=output_hidden_states,
     )
     if return_dict: