Update app.py

app.py

@@ -646,6 +646,14 @@ def train_model(data_path, epochs, batch_size, learning_rate, image_size, save_n
             
             avg_loss = epoch_loss / len(train_loader)
             logs.append(f"Epoch {epoch+1}/{epochs}: loss = {avg_loss:.4f}")
+
+            if (epoch + 1) % 10 == 0:
+                torch.save({
+                    "model_state_dict": MODEL.state_dict(),
+                    "text_encoder_state_dict": TEXT_ENCODER.state_dict(),
+                    "config": CONFIG
+                }, f"checkpoints/{save_name}_epoch{epoch+1}.pt")
+                logs.append(f"💾 Saved checkpoint at epoch {epoch+1}")
         
         # Save model
         MODEL.eval()
@@ -656,6 +664,19 @@ def train_model(data_path, epochs, batch_size, learning_rate, image_size, save_n
             "text_encoder_state_dict": TEXT_ENCODER.state_dict(),
             "config": CONFIG
         }, save_path)
+
+        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("☁️ Checkpoint uploaded to repo")
+        except Exception as e:
+            logs.append(f"⚠️ Upload failed: {e}")
         
         logs.append("-" * 40)
         logs.append(f"✅ Model saved to {save_path}")