Spaces:

TeamSAS
/

ImageColorizer

Sleeping

AUMREDKA commited on Aug 10

Commit

21d4284

verified ·

1 Parent(s): e697ec7

Test with checkpoint 24

Files changed (1) hide show

app.py CHANGED Viewed

@@ -17,8 +17,8 @@ from skimage.metrics import structural_similarity as ssim_metric
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 from model import ViTUNetColorizer
-CKPT = "checkpoints/checkpoint_epoch_015_20250808_154437.pt"
 model = None
 if os.path.exists(CKPT):
     model = ViTUNetColorizer(vit_model_name="vit_tiny_patch16_224").to(device)

 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 from model import ViTUNetColorizer
+# CKPT = "checkpoints/checkpoint_epoch_015_20250808_154437.pt"
+CKPT = "checkpoints/checkpoint_epoch_024_20250808_200347.pt"
 model = None
 if os.path.exists(CKPT):
     model = ViTUNetColorizer(vit_model_name="vit_tiny_patch16_224").to(device)