Upload app.py

app.py

@@ -390,6 +390,42 @@ DEVICE = None
 CONFIG = None
 
 
+def save_to_hub(save_name, repo_id=None):
+    """Save model checkpoint to HuggingFace Hub for persistence."""
+    global MODEL, TEXT_ENCODER, CONFIG
+    
+    if MODEL is None:
+        return "❌ No model loaded to save"
+    
+    try:
+        from huggingface_hub import HfApi, upload_file
+        import tempfile
+        
+        # Save to temp file
+        with tempfile.NamedTemporaryFile(suffix='.pt', delete=False) as f:
+            torch.save({
+                "model_state_dict": MODEL.state_dict(),
+                "text_encoder_state_dict": TEXT_ENCODER.state_dict(),
+                "config": CONFIG
+            }, f.name)
+            temp_path = f.name
+        
+        # Upload to Hub (same Space repo)
+        api = HfApi()
+        api.upload_file(
+            path_or_fileobj=temp_path,
+            path_in_repo=f"checkpoints/{save_name}.pt",
+            repo_id=repo_id or "Spanicin/candlestick-diffusion",
+            repo_type="space"
+        )
+        
+        os.unlink(temp_path)
+        return f"✅ Model saved to Hub: checkpoints/{save_name}.pt"
+    
+    except Exception as e:
+        return f"❌ Failed to save to Hub: {str(e)}"
+
+
 def load_model(checkpoint_path=None):
     global MODEL, TEXT_ENCODER, DIFFUSION, DEVICE, CONFIG
     
@@ -686,6 +722,21 @@ def train_model(data_path, epochs, batch_size, learning_rate, image_size, save_n
         logs.append("-" * 40)
         logs.append(f"✅ Model saved to {save_path}")
         
+        # Also try to save to Hub for persistence
+        try:
+            from huggingface_hub import HfApi
+            api = HfApi()
+            api.upload_file(
+                path_or_fileobj=save_path,
+                path_in_repo=f"checkpoints/{save_name}.pt",
+                repo_id="Spanicin/candlestick-diffusion",
+                repo_type="space"
+            )
+            logs.append(f"☁️ Model uploaded to Hub (persistent)")
+        except Exception as hub_error:
+            logs.append(f"⚠️ Could not upload to Hub: {hub_error}")
+            logs.append("   Model saved locally but may be lost on restart")
+        
         return "\n".join(logs)
     
     except Exception as e: