Fix MPI library issue

app.py

@@ -31,7 +31,7 @@ from train_settings.dvd.eval_utils import extract_raw_features_single,extract_ra
 from datasets.utils.warping import register_model2
 
 import gradio as gr
-
+from huggingface_hub import hf_hub_download
 
 
 
@@ -515,6 +515,7 @@ model, diffusion = create_model_and_diffusion(
 setattr(diffusion, "settings", settings)
 
 pretrained_dewarp_model = GeoTr_Seg_Inf()
+settings.env.seg_model_path = hf_hub_download(repo_id="hanquansanren/dvd", filename="seg.pth")
 reload_segmodel(pretrained_dewarp_model.msk, settings.env.seg_model_path)
 # reload_model(pretrained_dewarp_model.GeoTr, settings.env.dewarping_model_path)
 pretrained_dewarp_model.to(dist_util.dev())
@@ -523,16 +524,19 @@ pretrained_dewarp_model.eval()
 if settings.env.use_line_mask:
     pretrained_line_seg_model = UNet(n_channels=3, n_classes=1)
     pretrained_seg_model = Seg()
+    settings.env.line_seg_model_path = hf_hub_download(repo_id="hanquansanren/dvd", filename="line_model2.pth")
     line_model_ckpt = dist_util.load_state_dict(settings.env.line_seg_model_path, map_location='cpu')['model']
     pretrained_line_seg_model.load_state_dict(line_model_ckpt, strict=True)    
     pretrained_line_seg_model.to(dist_util.dev())
     pretrained_line_seg_model.eval()
 
+    settings.env.new_seg_model_path = hf_hub_download(repo_id="hanquansanren/dvd", filename="seg_model.pth")
     seg_model_ckpt = dist_util.load_state_dict(settings.env.new_seg_model_path, map_location='cpu')['model']
     pretrained_seg_model.load_state_dict(seg_model_ckpt, strict=True)    
     pretrained_seg_model.to(dist_util.dev())
     pretrained_seg_model.eval()
 
+settings.env.model_path = hf_hub_download(repo_id="hanquansanren/dvd", filename="model1852000.pt")
 model.cpu().load_state_dict(dist_util.load_state_dict(settings.env.model_path, map_location="cpu"), strict=False)
 logger.log(f"Model loaded with {settings.env.model_path}")
 

requirements.txt

@@ -3,7 +3,10 @@
 # pip install torch==2.1.1 torchvision==0.16.1 --index-url https://download.pytorch.org/whl/cu118
 # pip install mmcv==2.2.0 -f https://download.openmmlab.com/mmcv/dist/cu118/torch2.1/index.html
 # gxx_linux-64
+mmcv==2.2.0
 mpi4py
+matplotlib
+huggingface_hub
 opencv-python
 jpeg4py
 packaging