update two buttons

Video-P2P/run_videop2p.py

@@ -104,8 +104,8 @@ def main(
     mask_th = (.3, .3)
 
 
-    pretrained_model_path = pretrained_model_path
-    # pretrained_model_path = output_dir
+    # pretrained_model_path = pretrained_model_path
+    pretrained_model_path = output_dir
     image_path = train_data['video_path']
     prompt = train_data['prompt']
     # prompts = [prompt, ]
@@ -142,8 +142,8 @@ def main(
         pretrained_model_path,
         subfolder="vae",
     ).to(device, dtype=weight_dtype)
-    unet = UNet3DConditionModel.from_pretrained_2d(
-    # unet = UNet3DConditionModel.from_pretrained(
+    # unet = UNet3DConditionModel.from_pretrained_2d(
+    unet = UNet3DConditionModel.from_pretrained(
         pretrained_model_path, subfolder="unet"
     ).to(device)
     ldm_stable = TuneAVideoPipeline(

app_training.py

@@ -142,6 +142,18 @@ def create_training_demo(trainer: Trainer,
                 remove_gpu_after_training, input_token, blend_word_1, blend_word_2, eq_params_1, eq_params_2
             ],
             outputs=output_message)
+        run_button.click(
+            fn=trainer.run_p2p,
+            inputs=[
+                training_video, training_prompt, output_model_name,
+                delete_existing_repo, validation_prompt, base_model,
+                resolution, num_training_steps, learning_rate,
+                gradient_accumulation, seed, fp16, use_8bit_adam,
+                checkpointing_steps, validation_epochs, upload_to_hub,
+                use_private_repo, delete_existing_repo, upload_to,
+                remove_gpu_after_training, input_token, blend_word_1, blend_word_2, eq_params_1, eq_params_2
+            ],
+            outputs=output_message)
     return demo
 
 

trainer.py

@@ -104,6 +104,136 @@ class Trainer:
             shutil.rmtree(output_dir, ignore_errors=True)
         output_dir.mkdir(parents=True)
 
+        if upload_to_hub:
+            self.join_model_library_org(
+                self.hf_token if self.hf_token else input_token)
+
+        config = OmegaConf.load('Video-P2P/configs/man-skiing.yaml')
+        config.pretrained_model_path = self.download_base_model(base_model)
+        config.output_dir = output_dir.as_posix()
+        config.train_data.video_path = training_video.name  # type: ignore
+        config.train_data.prompt = training_prompt
+        config.train_data.n_sample_frames = 8
+        config.train_data.width = resolution
+        config.train_data.height = resolution
+        config.train_data.sample_start_idx = 0
+        config.train_data.sample_frame_rate = 1
+        config.validation_data.prompts = [validation_prompt]
+        config.validation_data.video_length = 8
+        config.validation_data.width = resolution
+        config.validation_data.height = resolution
+        config.validation_data.num_inference_steps = 50
+        config.validation_data.guidance_scale = 7.5
+        config.learning_rate = learning_rate
+        config.gradient_accumulation_steps = gradient_accumulation
+        config.train_batch_size = 1
+        config.max_train_steps = n_steps
+        config.checkpointing_steps = checkpointing_steps
+        config.validation_steps = validation_epochs
+        config.seed = seed
+        config.mixed_precision = 'fp16' if fp16 else ''
+        config.use_8bit_adam = use_8bit_adam
+        config.prompts = [training_prompt, validation_prompt]
+        config.blend_word = [blend_word_1, blend_word_2]
+        config.eq_params = {"words":[eq_params_1], "values":[int(eq_params_2)]}
+        if len(validation_prompt) == len(training_prompt):
+            config.is_word_swap = True
+        else:
+            config.is_word_swap = False
+
+        config_path = output_dir / 'config.yaml'
+        with open(config_path, 'w') as f:
+            OmegaConf.save(config, f)
+
+        command = f'accelerate launch Video-P2P/run_tuning.py --config {config_path}'
+        subprocess.run(shlex.split(command))
+        # command = f'python Video-P2P/run_videop2p.py --config {config_path}'
+        # subprocess.run(shlex.split(command))
+        save_model_card(save_dir=output_dir,
+                        base_model=base_model,
+                        training_prompt=training_prompt,
+                        test_prompt=validation_prompt,
+                        test_image_dir='results')
+
+        message = 'Training completed!'
+        print(message)
+
+        if upload_to_hub:
+            upload_message = self.model_uploader.upload_model(
+                folder_path=output_dir.as_posix(),
+                repo_name=output_model_name,
+                upload_to=upload_to,
+                private=use_private_repo,
+                delete_existing_repo=delete_existing_repo,
+                input_token=input_token)
+            print(upload_message)
+            message = message + '\n' + upload_message
+
+        if remove_gpu_after_training:
+            space_id = os.getenv('SPACE_ID')
+            if space_id:
+                api = HfApi(
+                    token=self.hf_token if self.hf_token else input_token)
+                api.request_space_hardware(repo_id=space_id,
+                                           hardware='cpu-basic')
+
+        return message
+
+
+    def run_p2p(
+        self,
+        training_video: str,
+        training_prompt: str,
+        output_model_name: str,
+        overwrite_existing_model: bool,
+        validation_prompt: str,
+        base_model: str,
+        resolution_s: str,
+        n_steps: int,
+        learning_rate: float,
+        gradient_accumulation: int,
+        seed: int,
+        fp16: bool,
+        use_8bit_adam: bool,
+        checkpointing_steps: int,
+        validation_epochs: int,
+        upload_to_hub: bool,
+        use_private_repo: bool,
+        delete_existing_repo: bool,
+        upload_to: str,
+        remove_gpu_after_training: bool,
+        input_token: str,
+        blend_word_1: str,
+        blend_word_2: str,
+        eq_params_1: str,
+        eq_params_2: str,
+    ) -> str:
+        # if SPACE_ID == ORIGINAL_SPACE_ID:
+        #     raise gr.Error(
+        #         'This Space does not work on this Shared UI. Duplicate the Space and attribute a GPU'
+        #     )
+        if not torch.cuda.is_available():
+            raise gr.Error('CUDA is not available.')
+        if training_video is None:
+            raise gr.Error('You need to upload a video.')
+        if not training_prompt:
+            raise gr.Error('The training prompt is missing.')
+        if not validation_prompt:
+            raise gr.Error('The validation prompt is missing.')
+
+        resolution = int(resolution_s)
+
+        if not output_model_name:
+            timestamp = datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
+            output_model_name = f'video-p2p-{timestamp}'
+        output_model_name = slugify.slugify(output_model_name)
+
+        repo_dir = pathlib.Path(__file__).parent
+        output_dir = repo_dir / 'experiments' / output_model_name
+        if overwrite_existing_model or upload_to_hub:
+            shutil.rmtree(output_dir, ignore_errors=True)
+        output_dir.mkdir(parents=True)
+
         if upload_to_hub:
             self.join_model_library_org(
                 self.hf_token if self.hf_token else input_token)
@@ -147,7 +277,6 @@ class Trainer:
 
         # command = f'accelerate launch Video-P2P/run_tuning.py --config {config_path}'
         # subprocess.run(shlex.split(command))
-        # torch.cuda.empty_cache()
         command = f'python Video-P2P/run_videop2p.py --config {config_path}'
         subprocess.run(shlex.split(command))
         save_model_card(save_dir=output_dir,