Add progress bar

requirements.txt

@@ -3,4 +3,5 @@ numpy>=2.0.1
 pillow>=10.4.0
 transformers>=4.44.0
 sentencepiece>=0.2.0
-torch>=2.3.1
+torch>=2.3.1
+tqdm>=4.66.5

src/flux_1_schnell/models/flux.py

@@ -1,6 +1,7 @@
 import PIL
 import mlx.core as mx
 from PIL import Image
+from tqdm import tqdm
 
 from flux_1_schnell.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
 from flux_1_schnell.tokenizer.clip_tokenizer import TokenizerCLIP
@@ -41,7 +42,7 @@ class Flux1Schnell:
             t5_text_encoder=self.t5_text_encoder
         )
 
-        for t in range(config.num_inference_steps):
+        for t in tqdm(range(config.num_inference_steps)):
             noise = self.transformer.predict(
                 t=t,
                 prompt_embeds=prompt_embeds,
@@ -57,6 +58,8 @@ class Flux1Schnell:
                 config=config
             )
 
+            mx.eval(latents)
+
         latents = Flux1Schnell._unpack_latents(latents)
         decoded = self.vae.decode(latents)
         return ImageUtil.to_image(decoded)

src/flux_1_schnell/models/transformer/transformer.py

@@ -64,8 +64,6 @@ class Transformer(nn.Module):
         hidden_states = self.norm_out.forward(hidden_states, text_embeddings)
         hidden_states = self.proj_out(hidden_states)
         noise = hidden_states
-        mx.eval(noise)
-        print(t)
         return noise
 
     @staticmethod