Spaces:
Paused
Paused
| import torch | |
| import numpy as np | |
| from omegaconf import DictConfig | |
| from optimize_utils import MultiTrajectory | |
| from stream_inference_wrapper import StreamInferenceWrapper | |
| def _extract_block_trajectories( | |
| multi_traj: MultiTrajectory, | |
| ) -> tuple[ | |
| list[list[dict[str, bool | list[torch.Tensor]]]], | |
| list[np.ndarray], | |
| np.ndarray | None, | |
| ]: | |
| """Extract block_trajectories from a MultiTrajectory in the format expected by begin_optimize. | |
| Returns: | |
| block_trajectories: block_num x N x trajectory dict | |
| Each trajectory dict has keys 'is_rotation', 'deltas', 'start_point', | |
| and optionally 'rotation_center'. | |
| masks: list of N masks corresponding to each trajectory | |
| movable_mask: the movable area mask for the whole image | |
| """ | |
| if multi_traj.trajectories is None or len(multi_traj.trajectories) == 0: | |
| return [], [], None | |
| movable_mask = multi_traj.movable_mask | |
| # Collect per-trajectory masks | |
| masks = [traj.mask for traj in multi_traj.trajectories] | |
| # Find the maximum number of blocks across all trajectories | |
| max_blocks = ( | |
| max( | |
| len(traj.block_trajectories) | |
| for traj in multi_traj.trajectories | |
| if traj.block_trajectories | |
| ) | |
| if any(traj.block_trajectories for traj in multi_traj.trajectories) | |
| else 0 | |
| ) | |
| if max_blocks == 0: | |
| return [], masks, movable_mask | |
| block_trajectories = [] | |
| for block_idx in range(max_blocks): | |
| block = [] | |
| for traj in multi_traj.trajectories: | |
| if traj.block_trajectories and block_idx < len(traj.block_trajectories): | |
| block.append(traj.block_trajectories[block_idx]) | |
| else: | |
| # Provide an empty placeholder | |
| block.append( | |
| { | |
| "is_rotation": False, | |
| "deltas": [], | |
| "start_point": (0, 0), | |
| } | |
| ) | |
| block_trajectories.append(block) | |
| # Assert: the N of each block in block_trajectories should equal the length of masks | |
| for block_idx, block in enumerate(block_trajectories): | |
| assert len(block) == len(masks), ( | |
| f"Block {block_idx} has {len(block)} trajectories, " f"but there are {len(masks)} masks" | |
| ) | |
| assert ((len(block_trajectories) == 0) and (movable_mask is None)) or ( | |
| (len(block_trajectories) > 0) and (movable_mask is not None) | |
| ), "block_trajectories and movable_mask must both be present or both be absent" | |
| return block_trajectories, masks, movable_mask | |
| class StreamDragInferenceWrapper(StreamInferenceWrapper): | |
| def __init__( | |
| self, | |
| stream_model_config: DictConfig, | |
| checkpoint_path: str, | |
| total_generate_block_number: int, | |
| use_ema: bool = True, | |
| seed: int = 0, | |
| ): | |
| super().__init__( | |
| stream_model_config=stream_model_config, | |
| checkpoint_path=checkpoint_path, | |
| total_generate_block_number=total_generate_block_number, | |
| use_ema=use_ema, | |
| seed=seed, | |
| ) | |
| self.previous_record_feature_list = None | |
| def inference( | |
| self, | |
| start_block_index: int, | |
| end_block_index: int, | |
| prompt: str, | |
| # below are for drag optimization | |
| multiple_trajectory: MultiTrajectory = None, | |
| ): | |
| assert start_block_index >= 0 | |
| assert end_block_index > start_block_index | |
| print(f""" | |
| {self.__class__.__name__}.inference(): | |
| {start_block_index = } | {end_block_index = } | |
| """) | |
| sampled_noise = self.get_sampled_noise(start_block_index, end_block_index) | |
| prompts = [prompt] | |
| # Extract block_trajectories, masks, and movable_mask from multiple_trajectory | |
| drag_optimize_target_latent_index = -1 | |
| if multiple_trajectory is not None: | |
| block_trajectories, masks, movable_mask = _extract_block_trajectories( | |
| multiple_trajectory | |
| ) | |
| assert multiple_trajectory.drag_or_animation_select in [ | |
| "Drag", | |
| "Animation", | |
| ] | |
| if multiple_trajectory.drag_or_animation_select == "Drag": | |
| drag_optimize_target_latent_index = 2 | |
| else: | |
| block_trajectories, masks, movable_mask = [], [], None | |
| if len(block_trajectories) > 0: | |
| is_drag_optimize = True | |
| else: | |
| is_drag_optimize = False | |
| initial_latents = self.get_initial_latents( | |
| start_block_index, | |
| ) | |
| if initial_latents is not None: | |
| print(f"{initial_latents.shape = }") | |
| print(f"{block_trajectories = }") | |
| print(f"{len(masks) = }") | |
| latents_result = self.pipeline.inference( | |
| noise=sampled_noise, | |
| text_prompts=prompts, | |
| return_latents=True, | |
| initial_latent=initial_latents, | |
| do_not_decode_video=True, | |
| do_not_recompute_initial_latents=True, | |
| # below are for drag optimization | |
| model_config=self.stream_model_config, | |
| previous_record_feature_list=self.previous_record_feature_list, | |
| is_drag_optimize=is_drag_optimize, | |
| block_trajectories=block_trajectories, | |
| masks=masks, | |
| movable_mask=movable_mask, | |
| drag_optimize_target_latent_index=drag_optimize_target_latent_index, | |
| ) | |
| if self.stream_model_config.drag_optim_config.record_feature_block_indexes: | |
| latents, record_attention_values_list = latents_result | |
| else: | |
| latents = latents_result | |
| record_attention_values_list = None | |
| if self.recorded_latents is None: | |
| self.recorded_latents = latents | |
| else: | |
| self.recorded_latents = torch.concat( | |
| [ | |
| self.recorded_latents[:, :0], | |
| latents, | |
| ], | |
| dim=1, | |
| ) | |
| if record_attention_values_list is not None: | |
| def dict_first_value(d: dict): | |
| return next(iter(d.values())) | |
| print(f"{record_attention_values_list.keys() = }") # denoising timesteps | |
| print( | |
| f"{dict_first_value(record_attention_values_list).keys() = }" | |
| ) # attention block layers | |
| print( | |
| f"{dict_first_value(dict_first_value(record_attention_values_list)).keys() = }" | |
| ) # attention types name | |
| print( | |
| f"{dict_first_value(dict_first_value(dict_first_value(record_attention_values_list))).shape = }" | |
| ) # [1, 3, 30, 52, 1536] | |
| else: | |
| print(f"{record_attention_values_list = }") | |
| self.previous_record_feature_list = record_attention_values_list | |
| self.decode_and_update_video(start_block_index, end_block_index) | |
| return ( | |
| self.video, | |
| self.video[self.block_to_video_index(start_block_index) :], | |
| ) | |
| def reset( | |
| self, | |
| ): | |
| super().reset() | |
| self.previous_record_feature_list = None | |