Update app.py

app.py

@@ -7,8 +7,9 @@ from datasets import load_dataset
 from huggingface_hub import Repository
 from huggingface_hub import HfApi, HfFolder, Repository, create_repo
 import os
-token = os.getenv('HF_WRITE')
-print(token)
+
+token = os.getenv('NEW_TOKEN')
+
 import gradio as gr
 from PIL import Image
 import os
@@ -23,7 +24,7 @@ big = False if device == torch.device('cpu') else True
 
 # Parameters
 IMG_SIZE = 1024 if big else 256
-BATCH_SIZE = 16 if big else 1
+BATCH_SIZE = 16 if big else 4
 EPOCHS = 12
 LR = 0.0002
 dataset_id = "K00B404/pix2pix_flux_set"
@@ -62,7 +63,7 @@ class UNetWrapper:
     def __init__(self, unet_model, repo_id):
         self.model = unet_model
         self.repo_id = repo_id
-        self.token = token  # Make sure this environment variable is set
+        self.token = os.getenv('NEW_TOKEN') # Make sure this environment variable is set
         self.api = HfApi()
 
     def push_to_hub(self):