Update app.py

app.py

@@ -17,11 +17,11 @@ from CLIP import load as load_clip
 
 # Device configuration
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-big = True if device == torch.device('cpu') else False
+big = False if device == torch.device('cpu') else True
 
 # Parameters
 IMG_SIZE = 1024 if big else 256
-BATCH_SIZE = 1 if big else 4
+BATCH_SIZE = 1 if big else 1
 EPOCHS = 12
 LR = 0.0002
 dataset_id = "K00B404/pix2pix_flux_set"
@@ -237,8 +237,8 @@ pipeline_tag: image-to-image
 
 ## Model Description
 Custom UNet model for Pix2Pix image translation.
-- **Image Size:** 1024
-- **Model Type:** Big (1024)
+- **Image Size:** {save_dict['model_config']['img_size']}
+- **Model Type:** {"big" if big else "small"}_UNet ({save_dict['model_config']['img_size']})
 
 ## Usage