Improve checkpoint caching check and set logdir for easier output finding

app.py

@@ -51,15 +51,24 @@ def get_checkpoint_path(checkpoint_file, default_local_path, repo_id=None, filen
     # Try downloading from Model Hub if configured
     if repo_id and filename and HF_HUB_AVAILABLE:
         try:
-            print(f"📥 Downloading checkpoint from {repo_id}...")
             # Use /workspace if available (persistent), otherwise /tmp
             cache_dir = "/workspace/checkpoints" if os.path.exists("/workspace") else "/tmp/checkpoints"
             os.makedirs(cache_dir, exist_ok=True)
             
-            # Add timeout and better error handling
+            # Check if already downloaded
+            possible_path = os.path.join(cache_dir, "models--" + repo_id.replace("/", "--"), "snapshots", "*", filename)
+            import glob
+            existing = glob.glob(possible_path)
+            if existing:
+                checkpoint_path = existing[0]
+                print(f"✅ Using cached checkpoint: {checkpoint_path}")
+                return checkpoint_path
+            
+            # Download with progress tracking
             import time
             start_time = time.time()
-            print(f"Starting download from {repo_id}...")
+            print(f"📥 Downloading checkpoint from {repo_id}...")
+            print(f"File size: ~15.5 GB - This may take 10-30 minutes")
             
             checkpoint_path = hf_hub_download(
                 repo_id=repo_id,
@@ -67,12 +76,11 @@ def get_checkpoint_path(checkpoint_file, default_local_path, repo_id=None, filen
                 cache_dir=cache_dir,
                 local_files_only=False,
                 resume_download=True,  # Resume if interrupted
-                timeout=600  # 10 minute timeout per request
             )
             
             elapsed = time.time() - start_time
-            print(f"Download completed in {elapsed:.1f} seconds")
-            print(f"✅ Checkpoint downloaded to: {checkpoint_path}")
+            print(f"✅ Download completed in {elapsed/60:.1f} minutes")
+            print(f"✅ Checkpoint at: {checkpoint_path}")
             return checkpoint_path
         except Exception as e:
             error_detail = str(e)
@@ -124,6 +132,7 @@ def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path
         output_dir.mkdir(exist_ok=True)
         
         # Run sampling command
+        # Set logdir to outputs directory for easier file finding
         cmd = [
             "python", "sample.py",
             "--model_config_path", config_path,
@@ -137,7 +146,8 @@ def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path
             "--finetuned_vae", "none",
             "--jacobi", "1",
             "--jacobi_th", "0.001",
-            "--jacobi_block_size", "16"
+            "--jacobi_block_size", "16",
+            "--logdir", str(output_dir)  # Set logdir to outputs
         ]
         
         status_msg += "Running generation...\n"