fix model shape error

model.py

@@ -3,7 +3,6 @@ import numpy as np
 import torch
 
 from torch import Tensor, nn
-
 from layers import (DoubleStreamBlock, EmbedND, LastLayer,
                                  MLPEmbedder, SingleStreamBlock,
                                  timestep_embedding)
@@ -11,11 +10,6 @@ from layers import (DoubleStreamBlock, EmbedND, LastLayer,
 import torch.distributed as dist
 from diffusers.models.embeddings import get_1d_sincos_pos_embed_from_grid
 
-from accelerate.logging import get_logger
-logger = get_logger(__name__, log_level="INFO")
-
-
-
 
 @dataclass
 class FluxParams:
@@ -27,7 +21,7 @@ class FluxParams:
     num_heads: int
     depth: int
     depth_single_blocks: int
-    axes_dim: list[int]
+    axes_dim: list
     theta: int
     qkv_bias: bool
     guidance_embed: bool
@@ -162,6 +156,11 @@ class Flux(nn.Module):
         ip_scale: Tensor = 1.0, 
         return_intermediate: bool = False,
     ):
+        inputs = [img, img_ids, txt, txt_ids, timesteps, y]
+        for i, input in enumerate(inputs):
+            if input.shape[0] != 1:
+                inputs[i] = input.unsqueeze(0)
+        img, img_ids, txt, txt_ids, timestpes, y = inputs
 
         if return_intermediate:
             intermediate_double = []
@@ -271,4 +270,4 @@ class Flux(nn.Module):
         if return_intermediate:
             return img, intermediate_double, intermediate_single
         else:
-            return img
+            return img.squeeze()