Create pipeline.py

pipeline.py

@@ -0,0 +1,83 @@
+from diffusers import DiffusionPipeline
+from invert import Inverter
+from generate import Generator
+from utils import init_model, seed_everything, get_frame_ids
+
+class VidToMePipeline(DiffusionPipeline):
+    def __init__(self, device="cuda", sd_version="2.1", float_precision="fp16", height=512, width=512):
+        # this will initlize the core pipeline components
+        pipe, scheduler, model_key = init_model(device, sd_version, None, "none", float_precision)
+        self.pipe = pipe
+        self.scheduler = scheduler
+        self.model_key = model_key
+        self.device = device
+        self.sd_version = sd_version
+        self.float_precision = float_precision
+        self.height = height
+        self.width = width
+
+    def __call__(self, video_path=None, video_prompt=None, edit_prompt=None, 
+                 control_type="none", n_timesteps=50, guidance_scale=7.5, 
+                 negative_prompt="ugly, blurry, low res", frame_range=None, 
+                 use_lora=False, seed=123, local_merge_ratio=0.9, global_merge_ratio=0.8):
+        
+        # dynamic config built from user inputs
+        config = self._build_config(video_path, video_prompt, edit_prompt, control_type, 
+                                    n_timesteps, guidance_scale, negative_prompt, 
+                                    frame_range, use_lora, seed, local_merge_ratio, global_merge_ratio)
+        
+        # seed for reproducibility - change as you need
+        seed_everything(config['seed'])
+
+        # inversion stage
+        print("Start inversion!")
+        inversion = Inverter(self.pipe, self.scheduler, config)
+        inversion(config['input_path'], config['inversion']['save_path'])
+
+        # generation stage
+        print("Start generation!")
+        generator = Generator(self.pipe, self.scheduler, config)
+        frame_ids = get_frame_ids(config['generation']['frame_range'], None)
+        generator(config['input_path'], config['generation']['latents_path'], 
+                  config['generation']['output_path'], frame_ids=frame_ids)
+        print(f"Output generated at: {config['generation']['output_path']}")
+
+    def _build_config(self, video_path, video_prompt, edit_prompt, control_type, 
+                      n_timesteps, guidance_scale, negative_prompt, frame_range, 
+                      use_lora, seed, local_merge_ratio, global_merge_ratio):
+        # constructing config dictionary from user prompts
+        config = {
+            'sd_version': self.sd_version,
+            'input_path': video_path,
+            'work_dir': "outputs/",
+            'height': self.height,
+            'width': self.width,
+            'inversion': {
+                'prompt': video_prompt or "Default video prompt.",
+                'save_path': "outputs/latents",
+                'steps': 50,
+                'save_intermediate': False
+            },
+            'generation': {
+                'control': control_type,
+                'guidance_scale': guidance_scale,
+                'n_timesteps': n_timesteps,
+                'negative_prompt': negative_prompt,
+                'prompt': edit_prompt or "Default edit prompt.",
+                'latents_path': "outputs/latents",
+                'output_path': "outputs/final",
+                'frame_range': frame_range or [0, 32],
+                'use_lora': use_lora,
+                'local_merge_ratio': local_merge_ratio,
+                'global_merge_ratio': global_merge_ratio
+            },
+            'seed': seed,
+            'device': self.device,
+            'float_precision': self.float_precision
+        }
+        return config
+
+# Sample usage
+pipeline = VidToMePipeline(device="cuda", sd_version="2.1", float_precision="fp16")
+pipeline(video_path="path/to/video.mp4", video_prompt="A beautiful scene of a sunset", 
+         edit_prompt="Make the sunset look more vibrant", control_type="depth", n_timesteps=50)