Fix bugs

app.py

@@ -46,6 +46,23 @@ snapshot_download(
 TEXT_ENCODER  = FrozenOpenCLIPEmbedder().eval()
 IMAGE_ENCODER = FrozenOpenCLIPImageEmbedderV2().eval()
 
+default_path = "checkpoints/LayerAnimate-Mix"
+scheduler = DDIMScheduler.from_pretrained(default_path, subfolder="scheduler")
+image_projector = Resampler.from_pretrained(default_path, subfolder="image_projector").eval()
+vae, vae_dualref = None, None
+if "I2V" or "Mix" in default_path:
+    vae           = AutoencoderKL.from_pretrained(default_path, subfolder="vae").eval()
+if "Interp" or "Mix" in default_path:
+    vae_dualref   = AutoencoderKL_Dualref.from_pretrained(default_path, subfolder="vae_dualref").eval()
+unet              = UNetModel.from_pretrained(default_path, subfolder="unet").eval()
+layer_controlnet  = LayerControlNet.from_pretrained(default_path, subfolder="layer_controlnet").eval()
+PIPELINE = AnimationPipeline(
+    vae=vae, vae_dualref=vae_dualref, text_encoder=TEXT_ENCODER, image_encoder=IMAGE_ENCODER, image_projector=image_projector,
+    unet=unet, layer_controlnet=layer_controlnet, scheduler=scheduler
+).to(device=DEVICE, dtype=WEIGHT_DTYPE)
+if "Interp" or "Mix" in default_path:
+    PIPELINE.vae_dualref.decoder.to(dtype=torch.float32)
+
 TRANSFORMS = transforms.Compose([
     transforms.Resize(min(HEIGHT, WIDTH)),
     transforms.CenterCrop((HEIGHT, WIDTH)),
@@ -72,17 +89,15 @@ def set_model(pretrained_model_path):
         vae_dualref   = AutoencoderKL_Dualref.from_pretrained(pretrained_model_path, subfolder="vae_dualref").eval()
     unet              = UNetModel.from_pretrained(pretrained_model_path, subfolder="unet").eval()
     layer_controlnet  = LayerControlNet.from_pretrained(pretrained_model_path, subfolder="layer_controlnet").eval()
-
-    PIPELINE = AnimationPipeline(
+    PIPELINE.update(
         vae=vae, vae_dualref=vae_dualref, text_encoder=TEXT_ENCODER, image_encoder=IMAGE_ENCODER, image_projector=image_projector,
         unet=unet, layer_controlnet=layer_controlnet, scheduler=scheduler
-    ).to(device=DEVICE, dtype=WEIGHT_DTYPE)
+    )
+    PIPELINE.to(device=DEVICE, dtype=WEIGHT_DTYPE)
     if "Interp" or "Mix" in pretrained_model_path:
         PIPELINE.vae_dualref.decoder.to(dtype=torch.float32)
     return pretrained_model_path
 
-set_model("checkpoints/LayerAnimate-Mix")
-
 def upload_image(image):
     image = TRANSFORMS(image)
     return image
@@ -92,7 +107,7 @@ def run(input_image, input_image_end, pretrained_model_path, seed,
         prompt, n_prompt, num_inference_steps, guidance_scale,
         *layer_args):
     generator = set_seed(seed, DEVICE)
-    global layer_tracking_points, PIPELINE
+    global layer_tracking_points
     args_layer_tracking_points = [layer_tracking_points[i].value for i in range(LAYER_CAPACITY)]
 
     args_layer_masks = layer_args[:LAYER_CAPACITY]