Update app.py

app.py

@@ -22,16 +22,14 @@ from diffusers import DiffusionPipeline, AnimateDiffPipeline, MotionAdapter, Eul
 import jax
 import jax.numpy as jnp
 
-class MyModel(nn.Module):
+class Model(nn.Module):
     def __init__(self):
         super().__init__()
-        self.register_buffer('a', torch.ones(1, 1))
+        self.register_buffer('buffer', torch.ones(1, 1))
 
-    def forward(self, x, extend=False):
-        if extend:
-            new_tensor = torch.randn(1, 1)
-            self.a = torch.cat([self.a, new_tensor], dim=0)
-        return x
+    def forward(self, x):
+        new_tensor = torch.randn(1, 1)
+        self.buffer = torch.cat([self.buffer, new_tensor], dim=0)
 
 def forest_schnell():
     PIPE = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, token=os.getenv("hf_token")).to("cuda")
@@ -204,9 +202,7 @@ def main():
     global time
     global last_motion
     global base
-    global model
     
-    model = MyModel()
     last_motion=None
     fps=20
     time=16
@@ -221,11 +217,13 @@ def main():
 
     repo="stabilityai/sd-vae-ft-mse-original"
     ckpt="vae-ft-mse-840000-ema-pruned.safetensors"
-    vae = model(load_file(hf_hub_download(repo, ckpt), device=device))
+    vae = Model()
+    vae(load_file(hf_hub_download(repo, ckpt), device=device))
 
     repo="ByteDance/SDXL-Lightning"
     ckpt=f"sdxl_lightning_{step}step_unet.safetensors"
-    unet = model(load_file(hf_hub_download(repo, ckpt), device=device))
+    unet = Model()
+    unet(load_file(hf_hub_download(repo, ckpt), device=device))
     
     #repo = "SG161222/Realistic_Vision_V6.0_B1_noVAE"