Add better progress tracking and timing information for checkpoint download

app.py

@@ -83,14 +83,30 @@ def get_checkpoint_path(checkpoint_file, default_local_path, repo_id=None, filen
             start_time = time.time()
             print(f"📥 Downloading checkpoint from {repo_id}...")
             print(f"File size: ~15.5 GB - This may take 10-30 minutes")
+            print(f"Cache directory: {cache_dir}")
+            print(f"Progress will be shown below...")
             
-            checkpoint_path = hf_hub_download(
-                repo_id=repo_id,
-                filename=filename,
-                cache_dir=cache_dir,
-                local_files_only=False,
-                resume_download=True,  # Resume if interrupted
-            )
+            # Use tqdm for progress if available
+            try:
+                from tqdm import tqdm
+                from huggingface_hub.utils import tqdm as hf_tqdm
+                checkpoint_path = hf_hub_download(
+                    repo_id=repo_id,
+                    filename=filename,
+                    cache_dir=cache_dir,
+                    local_files_only=False,
+                    resume_download=True,  # Resume if interrupted
+                )
+            except Exception as e:
+                # Fallback if tqdm fails
+                print(f"Note: Progress bar unavailable, downloading silently...")
+                checkpoint_path = hf_hub_download(
+                    repo_id=repo_id,
+                    filename=filename,
+                    cache_dir=cache_dir,
+                    local_files_only=False,
+                    resume_download=True,
+                )
             
             elapsed = time.time() - start_time
             print(f"✅ Download completed in {elapsed/60:.1f} minutes")
@@ -148,7 +164,11 @@ def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path
         return None, "Error: Config file not found. Please ensure config file exists."
     
     status_msg += "Starting image generation...\n"
-    status_msg += "This may take 1-3 minutes for first run (model loading).\n"
+    status_msg += "⏱️  Timing breakdown:\n"
+    status_msg += "   - Checkpoint download: 10-30 min (first time only, ~15.5 GB)\n"
+    status_msg += "   - Model loading: 2-5 min (first time only)\n"
+    status_msg += "   - Image generation: 1-3 min\n"
+    status_msg += "   - Subsequent runs: Only generation time (~1-3 min)\n"
     
     try:
         # Create output directory
@@ -173,8 +193,8 @@ def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path
             "--logdir", str(output_dir)  # Set logdir to outputs directory
         ]
         
-        status_msg += "Running generation...\n"
-        status_msg += "Note: First run includes checkpoint download (~10-20 min) and model loading (~2-5 min).\n"
+        status_msg += "🚀 Running generation...\n"
+        status_msg += "📊 Current step: Model inference (checkpoint should already be downloaded)\n"
         
         # Create log file for debugging
         log_file = output_dir / "generation.log"