update app

webUI.py → app.py

@@ -35,6 +35,34 @@ from annotator.hed import HEDdetector
 from annotator.canny import CannyDetector
 from annotator.midas import MidasDetector
 
+import huggingface_hub
+
+huggingface_hub.hf_hub_download('SingleZombie/FRESCO',
+                                'boxer-punching-towards-camera.mp4',
+                                local_dir='data')
+huggingface_hub.hf_hub_download('SingleZombie/FRESCO',
+                                'car-turn.mp4',
+                                local_dir='data')
+huggingface_hub.hf_hub_download('SingleZombie/FRESCO',
+                                'dog.mp4',
+                                local_dir='data')
+huggingface_hub.hf_hub_download('SingleZombie/FRESCO',
+                                'music.mp4',
+                                local_dir='data')
+
+huggingface_hub.hf_hub_download('PKUWilliamYang/Rerender',
+                                'gmflow_sintel-0c07dcb3.pth',
+                                local_dir='model')
+
+huggingface_hub.hf_hub_download('PKUWilliamYang/Rerender',
+                                'epoch_resnet.pth',
+                                local_dir='model')
+
+huggingface_hub.hf_hub_download('PKUWilliamYang/Rerender',
+                                'ebsynth',
+                                local_dir='src/ebsynth/deps/ebsynth/bin')
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 def get_models(config):
     # optical flow
@@ -45,7 +73,7 @@ def get_models(config):
                         attention_type='swin',
                         ffn_dim_expansion=4,
                         num_transformer_layers=6,
-                        ).to('cuda')
+                        ).to(device)
 
     checkpoint = torch.load(
         config['gmflow_path'], map_location=lambda storage, loc: storage)
@@ -56,14 +84,14 @@ def get_models(config):
     # saliency detection
     sod_model = build_model('resnet')
     sod_model.load_state_dict(torch.load(config['sod_path']))
-    sod_model.to("cuda").eval()
+    sod_model.to(device).eval()
 
     # controlnet
     if config['controlnet_type'] not in ['hed', 'depth', 'canny']:
         config['controlnet_type'] = 'hed'
     controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-"+config['controlnet_type'],
                                                  torch_dtype=torch.float16)
-    controlnet.to("cuda")
+    controlnet.to(device)
     if config['controlnet_type'] == 'depth':
         detector = MidasDetector()
     elif config['controlnet_type'] == 'canny':
@@ -77,7 +105,7 @@ def get_models(config):
     pipe = StableDiffusionPipeline.from_pretrained(
         config['sd_path'], vae=vae, torch_dtype=torch.float16)
     pipe.scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
-    pipe.to("cuda")
+    pipe.to(device)
     pipe.scheduler.set_timesteps(
         config['num_inference_steps'], device=pipe._execution_device)
 
@@ -168,14 +196,15 @@ class GlobalState:
         self.control_type = control_type
         self.controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-"+control_type,
                                                           torch_dtype=torch.float16)
-        self.controlnet.to("cuda")
+        self.controlnet.to(device)
         if control_type == 'depth':
             self.detector = MidasDetector()
         elif control_type == 'canny':
             self.detector = CannyDetector()
         else:
             self.detector = HEDdetector()
-        torch.cuda.empty_cache()
+        if device == 'cuda':
+            torch.cuda.empty_cache()
         for param in self.controlnet.parameters():
             param.requires_grad = False
 
@@ -189,10 +218,11 @@ class GlobalState:
             sd_model, vae=vae, torch_dtype=torch.float16)
         self.pipe.scheduler = DDPMScheduler.from_config(
             self.pipe.scheduler.config)
-        self.pipe.to("cuda")
+        self.pipe.to(device)
         self.frescoProc = apply_FRESCO_attn(self.pipe)
         self.frescoProc.controller.disable_controller()
-        torch.cuda.empty_cache()
+        if device == 'cuda':
+            torch.cuda.empty_cache()
         for param in self.pipe.unet.parameters():
             param.requires_grad = False
 
@@ -291,7 +321,7 @@ def process1(input_path, prompt, sd_model, seed, image_resolution, control_stren
 
         edges = torch.cat([numpy2tensor(apply_control(img,
                                                       global_state.detector, control_type)[:, :, None]) for img in imgs], dim=0)
-        edges = edges.repeat(1, 3, 1, 1).cuda() * 0.5 + 0.5
+        edges = edges.repeat(1, 3, 1, 1).to(device) * 0.5 + 0.5
         edges = torch.cat([edges.to(global_state.pipe.unet.dtype)] * 2)
 
         if bg_smooth:
@@ -326,7 +356,8 @@ def process1(input_path, prompt, sd_model, seed, image_resolution, control_stren
                          saliency=saliency, optimize_temporal=optimize_temporal)
 
         gc.collect()
-        torch.cuda.empty_cache()
+        if device == 'cuda':
+            torch.cuda.empty_cache()
 
         # run!
         latents = inference(global_state.pipe, global_state.controlnet, global_state.frescoProc,
@@ -354,7 +385,8 @@ def process1(input_path, prompt, sd_model, seed, image_resolution, control_stren
         propagation_mode = batch_ind > 0
         if batch_ind == len(sublists):
             gc.collect()
-            torch.cuda.empty_cache()
+            if device == 'cuda':
+                torch.cuda.empty_cache()
             break
 
     writer = imageio.get_writer(os.path.join(save_path, 'key.mp4'), fps=fps)