Upload app.py

app.py

@@ -607,7 +607,7 @@ def train_model(data_path, epochs, batch_size, learning_rate, image_size, save_n
         train_dataset = ChartDataset(data_path, image_size=image_size, split="train")
         train_loader = DataLoader(
             train_dataset, batch_size=batch_size, shuffle=True,
-            num_workers=0, pin_memory=True, drop_last=True, collate_fn=collate_fn
+            num_workers=2, pin_memory=True, drop_last=True, collate_fn=collate_fn
         )
         
         # Optimizer
@@ -642,12 +642,9 @@ def train_model(data_path, epochs, batch_size, learning_rate, image_size, save_n
                 
                 epoch_loss += loss.item()
                 current_step += 1
-                print(f"Step {current_step}, loss: {loss.item():.4f}")
             
             avg_loss = epoch_loss / len(train_loader)
             logs.append(f"Epoch {epoch+1}/{epochs}: loss = {avg_loss:.4f}")
-
-        
         
         # Save model
         MODEL.eval()
@@ -658,19 +655,6 @@ 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}")