Update app.py

test on .from_pretrained

app.py

@@ -34,6 +34,7 @@ def load_dit_model(dit_size):
     # Configure model based on size
     if dit_size == "S":
         model = DiT(num_blocks=8, hidden_size=384, num_heads=6)
+        model.from_pretrained("kaupane/DiT-Wikiart-Small")
     elif dit_size == "B":
         model = DiT(num_blocks=12, hidden_size=640, num_heads=10)
     elif dit_size == "L":
@@ -42,8 +43,8 @@ def load_dit_model(dit_size):
         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"])
+    #checkpoint = torch.load(ckpt_path, map_location="cpu")
+    #model.load_state_dict(checkpoint["model_state_dict"])
     
     return model