Update app.py

app.py

@@ -31,17 +31,21 @@ from models.eva_clip import create_model_and_transforms
 from models.eva_clip.constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
 from models.eva_clip.utils_qformer import resize_numpy_image_long
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
 
-hf_hub_download(repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x4.pth", local_dir="model_real_esran")
-snapshot_download(repo_id="AlexWortega/RIFE", local_dir="model_rife")
-snapshot_download(repo_id="BestWishYsh/ConsisID-preview", local_dir="BestWishYsh/ConsisID-preview")
+model_path = "ckpts"
 
-model_path = "BestWishYsh/ConsisID-preview"
 lora_path = None
 lora_rank = 128
 dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
+if not os.path.exists(model_path) or not os.path.exists(f"{model_path}/model_real_esran") or not os.path.exists(f"{model_path}/model_rife"):
+    hf_hub_download(repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x4.pth", local_dir=f"{model_path}/model_real_esran")
+    snapshot_download(repo_id="AlexWortega/RIFE", local_dir=f"{model_path}/model_rife")
+    snapshot_download(repo_id="BestWishYsh/ConsisID-preview", local_dir=f"{model_path}")
+else:
+    print(f"Model already exists in {model_path}, skipping download.")
+    
 if os.path.exists(os.path.join(model_path, "transformer_ema")):
     subfolder = "transformer_ema"
 else:
@@ -110,8 +114,12 @@ if "variance_type" in pipe.scheduler.config:
 pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, **scheduler_args)
 pipe.to(device)
 
+# Enable CPU offload for the model.
+# turn on if you don't have multiple GPUs or enough GPU memory(such as H100) and it will cost more time in inference, it may also reduce the quality
 pipe.enable_model_cpu_offload()
 pipe.enable_sequential_cpu_offload()
+# pipe.vae.enable_slicing()
+# pipe.vae.enable_tiling()
 
 os.makedirs("./output", exist_ok=True)
 os.makedirs("./gradio_tmp", exist_ok=True)