Update app.py

app.py

@@ -648,12 +648,26 @@ def train_model(data_path, epochs, batch_size, learning_rate, image_size, save_n
             logs.append(f"Epoch {epoch+1}/{epochs}: loss = {avg_loss:.4f}")
 
             if (epoch + 1) % 10 == 0:
+                ckpt_path = f"checkpoints/{save_name}_epoch{epoch+1}.pt"
                 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")
+                }, ckpt_path)
                 logs.append(f"💾 Saved checkpoint at epoch {epoch+1}")
+                try:
+                    from huggingface_hub import HfApi
+                    api = HfApi()
+                    api.upload_file(
+                        path_or_fileobj=ckpt_path,
+                        path_in_repo=ckpt_path,
+                        repo_id="Spanicin/candlestick-diffusion",
+                        repo_type="space",
+                        token=os.environ.get("HF_TOKEN")
+                    )
+                    logs.append(f"☁️ Uploaded to repo")
+                except Exception as e:
+                    logs.append(f"⚠️ Upload failed: {e}")
         
         # Save model
         MODEL.eval()