Update app.py

app.py

@@ -15,8 +15,9 @@ from torch import multiprocessing as mp
 from pathos.threading import ThreadPool as Pool
 from diffusers.pipelines.flux import FluxPipeline
 from diffusers.utils import export_to_gif, load_image
+from diffusers.models.modeling_utils import ModelMixin
 from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file
+from safetensors.torch import load_file, save_file
 from diffusers import DiffusionPipeline, AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler, StableDiffusionXLPipeline, UNet2DConditionModel
 import jax
 import jax.numpy as jnp
@@ -207,12 +208,14 @@ def main():
 
     repo="stabilityai/sd-vae-ft-mse-original"
     ckpt="vae-ft-mse-840000-ema-pruned.safetensors"
-    vae = load_file(hf_hub_download(repo, ckpt), device=device)
+    save_file(load_file(hf_hub_download(repo, ckpt), device=device),"./vae")
+    vae = ModelMixin("./vae")
 
     repo="ByteDance/SDXL-Lightning"
     ckpt=f"sdxl_lightning_{step}step_unet.safetensors"
-    unet = load_file(hf_hub_download(repo, ckpt), device=device)
-
+    save_file(load_file(hf_hub_download(repo, ckpt), device=device),"./unet")
+    unet = ModelMixin("./unet")
+    
     #repo = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
     
     #repo = "emilianJR/epiCRealism"