Update app.py

app.py

@@ -23,13 +23,6 @@ USE_HALF_PRECISION = True
 
 def load_dit_model(dit_size):
     """Load DiT model of specified size"""
-    #ckpt_path = f"./ckpts/DiT_{dit_size}_final.pth"
-    ckpt_path = hf_hub_download(
-        repo_id = "kaupane/DiT-Wikiart",
-        filename = f"DiT_{dit_size}_final.pth"
-    )
-    if not os.path.exists(ckpt_path):
-        raise FileNotFoundError(f"Checkpoint not found at {ckpt_path}")
     
     # Configure model based on size
     if dit_size == "S":
@@ -37,15 +30,13 @@ def load_dit_model(dit_size):
         model.from_pretrained("kaupane/DiT-Wikiart-Small")
     elif dit_size == "B":
         model = DiT(num_blocks=12, hidden_size=640, num_heads=10)
+        model.from_pretrained("kaupane/DiT-Wikiart-Base")
     elif dit_size == "L":
         model = DiT(num_blocks=16, hidden_size=896, num_heads=14)
+        model.from_pretrained("kaupane/DiT-Wikiart-Large")
     else:
         raise ValueError(f"Invalid DiT size: {dit_size}")
     
-    # Load checkpoint
-    #checkpoint = torch.load(ckpt_path, map_location="cpu")
-    #model.load_state_dict(checkpoint["model_state_dict"])
-    
     return model
 
 class DiffusionSampler: