kaupane commited on
Commit
2f32e9a
·
verified ·
1 Parent(s): db88ac3

Update app.py

Browse files

Corrected a stupid mistake

Files changed (1) hide show
  1. app.py +3 -6
app.py CHANGED
@@ -26,14 +26,11 @@ def load_dit_model(dit_size):
26
 
27
  # Configure model based on size
28
  if dit_size == "S":
29
- model = DiT(num_blocks=8, hidden_size=384, num_heads=6)
30
- model.from_pretrained("kaupane/DiT-Wikiart-Small")
31
  elif dit_size == "B":
32
- model = DiT(num_blocks=12, hidden_size=640, num_heads=10)
33
- model.from_pretrained("kaupane/DiT-Wikiart-Base")
34
  elif dit_size == "L":
35
- model = DiT(num_blocks=16, hidden_size=896, num_heads=14)
36
- model.from_pretrained("kaupane/DiT-Wikiart-Large")
37
  else:
38
  raise ValueError(f"Invalid DiT size: {dit_size}")
39
 
 
26
 
27
  # Configure model based on size
28
  if dit_size == "S":
29
+ model = DiT.from_pretrained("kaupane/DiT-Wikiart-Small")
 
30
  elif dit_size == "B":
31
+ model = DiT.from_pretrained("kaupane/DiT-Wikiart-Base")
 
32
  elif dit_size == "L":
33
+ model = DiT.from_pretrained("kaupane/DiT-Wikiart-Large")
 
34
  else:
35
  raise ValueError(f"Invalid DiT size: {dit_size}")
36