Improve error messages and add status updates for long-running operations

app.py

@@ -60,12 +60,16 @@ def get_checkpoint_path(checkpoint_file, default_local_path, repo_id=None, filen
                 repo_id=repo_id,
                 filename=filename,
                 cache_dir=cache_dir,
-                local_files_only=False
+                local_files_only=False,
+                resume_download=True  # Resume if interrupted
             )
             print(f"✅ Checkpoint downloaded to: {checkpoint_path}")
             return checkpoint_path
         except Exception as e:
-            return None, f"Error downloading checkpoint: {str(e)}"
+            error_detail = str(e)
+            if "404" in error_detail or "not found" in error_detail.lower():
+                return None, f"Checkpoint not found in Model Hub.\n\nPlease verify:\n1. Repo exists: https://huggingface.co/{repo_id}\n2. File exists: {filename}\n3. Repo is Public (not Private)\n\nError: {error_detail}"
+            return None, f"Error downloading checkpoint: {error_detail}\n\nThis may take 10-30 minutes for a 14GB file. Please wait or check your internet connection."
     
     # No checkpoint found
     return None, f"Checkpoint not found. Please upload a checkpoint file or configure Model Hub repository."
@@ -92,12 +96,19 @@ def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path
         return None, result[1]  # Error message
     
     checkpoint_path = result
+    
+    # Show status
+    status_msg = f"Using checkpoint: {checkpoint_path}\n"
+    
     if not os.path.exists(checkpoint_path):
-        return None, f"Error: Checkpoint file not found at {checkpoint_path}."
+        return None, f"Error: Checkpoint file not found at {checkpoint_path}.\n\nPlease verify:\n1. Model Hub repo exists: {IMAGE_CHECKPOINT_REPO}\n2. File name matches: starflow_3B_t2i_256x256.pth\n3. Repo is Public (not Private)"
     
     if not config_path or not os.path.exists(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"
+    
     try:
         # Create output directory
         output_dir = Path("outputs")
@@ -120,10 +131,16 @@ def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path
             "--jacobi_block_size", "16"
         ]
         
-        result = subprocess.run(cmd, capture_output=True, text=True, cwd=os.getcwd())
+        status_msg += "Running generation...\n"
+        
+        # Run with timeout (30 minutes max)
+        result = subprocess.run(cmd, capture_output=True, text=True, cwd=os.getcwd(), timeout=1800)
         
         if result.returncode != 0:
-            return None, f"Error: {result.stderr}"
+            error_msg = f"Error during generation:\n{result.stderr}\n\nStdout:\n{result.stdout}"
+            return None, error_msg
+        
+        status_msg += "Generation complete. Looking for output...\n"
         
         # Find the generated image
         # The sample.py script saves to logdir/model_name/...
@@ -131,9 +148,9 @@ def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path
         output_files = list(output_dir.glob("**/*.png")) + list(output_dir.glob("**/*.jpg"))
         if output_files:
             latest_file = max(output_files, key=lambda p: p.stat().st_mtime)
-            return str(latest_file), "Success! Image generated."
+            return str(latest_file), status_msg + "✅ Success! Image generated."
         else:
-            return None, "Error: Generated image not found."
+            return None, status_msg + f"Error: Generated image not found in {output_dir}. Check stdout:\n{result.stdout}"
             
     except Exception as e:
         return None, f"Error: {str(e)}"