Update test_flux.py

test_flux.py

@@ -1,3 +1,5 @@
+#!/usr/bin/env python3
+
 import time
 import random
 import torch
@@ -6,6 +8,14 @@ from PIL import Image
 import sys
 import os
 
+# Add TotoroUI to path
+sys.path.append(os.getcwd())
+
+print("=== FLUX Inference Timing Test ===")
+print(f"CUDA available: {torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    print(f"GPU: {torch.cuda.get_device_name(0)}")
+    print(f"CUDA version: {torch.version.cuda}")
 
 import nodes
 from nodes import NODE_CLASS_MAPPINGS
@@ -25,6 +35,7 @@ def closestNumber(n, m):
 
 # Load models
 print("\n⏳ Loading models...")
+start_time = time.time()
 
 DualCLIPLoader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
 UNETLoader = NODE_CLASS_MAPPINGS["UNETLoader"]()
@@ -43,7 +54,8 @@ with torch.inference_mode():
     unet = UNETLoader.load_unet("flux1-dev-fp8.safetensors", "fp8_e4m3fn")[0]
     vae = VAELoader.load_vae("ae.sft")[0]
 
-
+load_time = time.time() - load_start
+print(f"✅ Models loaded in {load_time:.2f}s")
 
 # Inference test
 print("\n🚀 Starting inference...")
@@ -77,7 +89,22 @@ with torch.inference_mode():
     model_management.soft_empty_cache()
     
     decoded = VAEDecode.decode(vae, sample)[0].detach()
-    image = Image.fromarray(np.array(decoded*255, dtype=np.uint8)[0])
+    # Fix deprecation warning
+    image_array = (decoded*255).clamp(0, 255).byte().cpu().numpy()[0]
+    image = Image.fromarray(image_array)
     image.save("flux_output.png")
 
-inference_time = time.time() - inference_start
+inference_time = time.time() - inference_start
+total_time = time.time() - start_time
+
+print(f"\n⏱️  === TIMING RESULTS ===")
+print(f"🔄 Model loading: {load_time:.2f}s")
+print(f"🎨 Image generation: {inference_time:.2f}s") 
+print(f"⏱️  Total time: {total_time:.2f}s")
+print(f"💾 Image saved: flux_output.png")
+
+# Performance metrics
+steps_per_sec = steps / inference_time
+print(f"\n📊 Performance:")
+print(f"Steps/second: {steps_per_sec:.2f}")
+print(f"Time per step: {inference_time/steps:.2f}s")