Update app.py

app.py

@@ -165,7 +165,11 @@ def sample_weights(unet, proj, mean, std, v, device, factor = 1.0):
 @torch.no_grad()
 @spaces.GPU
 def sample_model():
-    unet.value, _, _, _, _ = load_models()
+    unet.value = UNet2DConditionModel.from_pretrained(
+        pretrained_model_name_or_path, subfolder="unet", revision=revision
+    )
+    unet.value.requires_grad_(False)
+    unet.value.to(device.value, dtype=weight_dtype)
     network.value = sample_weights(unet.value, proj, mean, std, v[:, :1000], device.value, factor = 1.00)
  
 @torch.no_grad()