Update app.py

app.py

@@ -219,11 +219,13 @@ def main():
     ckpt="vae-ft-mse-840000-ema-pruned.safetensors"
     vae = Model()
     vae(load_file(hf_hub_download(repo, ckpt), device=device))
+    vae = ModelMixin(vae)
 
     repo="ByteDance/SDXL-Lightning"
     ckpt=f"sdxl_lightning_{step}step_unet.safetensors"
     unet = Model()
     unet(load_file(hf_hub_download(repo, ckpt), device=device))
+    unet = ModelMixin(unet)
     
     #repo = "SG161222/Realistic_Vision_V6.0_B1_noVAE"