Update app.py

app.py

@@ -31,6 +31,12 @@ class Model(nn.Module):
         new_tensor = torch.randn(1, 1)
         self.buffer = torch.cat([self.buffer, new_tensor], dim=0)
 
+def dict2model(dict,model=Model()):
+    model(dict)
+    mix = ModelMixin()
+    mix(model)
+    return mix
+
 def forest_schnell():
     PIPE = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, token=os.getenv("hf_token")).to("cuda")
     return PIPE
@@ -217,15 +223,11 @@ def main():
 
     repo="stabilityai/sd-vae-ft-mse-original"
     ckpt="vae-ft-mse-840000-ema-pruned.safetensors"
-    vae = Model()
-    vae(load_file(hf_hub_download(repo, ckpt), device=device))
-    vae = ModelMixin(vae)
+    vae = dict2model(load_file(hf_hub_download(repo, ckpt), device=device))
 
     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)
+    unet = dict2model(load_file(hf_hub_download(repo, ckpt), device=device))
     
     #repo = "SG161222/Realistic_Vision_V6.0_B1_noVAE"