Update app.py

app.py

@@ -22,6 +22,21 @@ from diffusers import DiffusionPipeline, AnimateDiffPipeline, MotionAdapter, Eul
 import jax
 import jax.numpy as jnp
 
+class MyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer('a', torch.ones(1, 1))
+
+    def forward(self, x: torch.Tensor, extend: bool):
+        if extend:
+            new_tensor = torch.randn(1, 1)
+            self.a = torch.cat([self.a, new_tensor], dim=0)
+        return x
+
+for _ in range(10):
+    out = model(x, extend=True)
+    print(model.state_dict())
+
 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
@@ -193,7 +208,9 @@ def main():
     global time
     global last_motion
     global base
-
+    global model
+    
+    model = MyModel()
     last_motion=None
     fps=20
     time=16
@@ -208,13 +225,12 @@ def main():
 
     repo="stabilityai/sd-vae-ft-mse-original"
     ckpt="vae-ft-mse-840000-ema-pruned.safetensors"
-    save_file(load_file(hf_hub_download(repo, ckpt), device=device),"./vae")
+    vae = model(load_file(hf_hub_download(repo, ckpt), device=device)
     vae = "./vae"
 
     repo="ByteDance/SDXL-Lightning"
     ckpt=f"sdxl_lightning_{step}step_unet.safetensors"
-    save_file(load_file(hf_hub_download(repo, ckpt), device=device),"./unet")
-    unet = "./unet"
+    unet = model(load_file(hf_hub_download(repo, ckpt), device=device)
     
     #repo = "SG161222/Realistic_Vision_V6.0_B1_noVAE"