updates

app.py

@@ -12,10 +12,20 @@ from huggingface_hub import snapshot_download
 import os
 
 # Download and load the model
-model_path = 'pyramid_flow_model'
+model_path = os.path.join(os.getcwd(), 'pyramid_flow_model')
 if not os.path.exists(model_path):
     snapshot_download("rain1011/pyramid-flow-sd3", local_dir=model_path, local_dir_use_symlinks=False, repo_type='model')
 
+# Verify model files exist
+dit_path = os.path.join(model_path, 'diffusion_transformer_768p')
+vae_path = os.path.join(model_path, 'causal_video_vae')
+
+if not os.path.exists(os.path.join(dit_path, 'diffusion_pytorch_model.safetensors')):
+    raise FileNotFoundError(f"DiT model file not found in {dit_path}")
+
+if not os.path.exists(os.path.join(vae_path, 'diffusion_pytorch_model.safetensors')):
+    raise FileNotFoundError(f"VAE model file not found in {vae_path}")
+
 torch.cuda.set_device(0)
 model_dtype, torch_dtype = 'bf16', torch.bfloat16