Spaces:
Paused
Paused
File size: 5,100 Bytes
186aa49 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | import torch
from ...pipelines.cosmos.pipeline_cosmos3_omni import Cosmos3OmniPipeline, CosmosSafetyChecker
from ..modular_pipeline import ModularPipeline
class Cosmos3OmniModularPipeline(ModularPipeline):
"""
A ModularPipeline for Cosmos 3 omni generation.
"""
default_blocks_name = "Cosmos3OmniBlocks"
duration_template = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS."
image_resolution_template = "This image is of {height}x{width} resolution."
video_resolution_template = "This video is of {height}x{width} resolution."
inverse_duration_template = "The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS."
inverse_image_resolution_template = "This image is not of {height}x{width} resolution."
inverse_video_resolution_template = "This video is not of {height}x{width} resolution."
@property
def vae_scale_factor_spatial(self):
if getattr(self, "vae", None) is not None:
return int(self.vae.config.scale_factor_spatial)
return 16
@property
def vae_scale_factor_temporal(self):
if getattr(self, "vae", None) is not None:
return int(self.vae.config.scale_factor_temporal)
return 4
@property
def num_channels_latents(self):
if getattr(self, "transformer", None) is not None:
return int(self.transformer.config.latent_channel)
return 48
@property
def sound_sampling_rate(self):
if getattr(self, "sound_tokenizer", None) is not None:
return int(self.sound_tokenizer.config.sampling_rate)
return 48000
@property
def sound_hop_size(self):
if getattr(self, "sound_tokenizer", None) is not None:
return int(self.sound_tokenizer._hop_size)
return 1920
@property
def _vae_latents_mean(self):
return torch.tensor(self.vae.config.latents_mean, dtype=self.vae.dtype)
@property
def _vae_latents_inv_std(self):
return 1.0 / torch.tensor(self.vae.config.latents_std, dtype=self.vae.dtype)
@property
def llm_special_tokens(self):
if getattr(self, "text_tokenizer", None) is None:
return None
return {
"start_of_generation": self.text_tokenizer.convert_tokens_to_ids("<|vision_start|>"),
"eos_token_id": self.text_tokenizer.eos_token_id,
}
def enable_safety_checker(self, safety_checker=None):
if safety_checker is not None:
self.safety_checker = safety_checker
elif getattr(self, "safety_checker", None) is None:
self.safety_checker = CosmosSafetyChecker()
self._is_safety_checker_enabled = True
def disable_safety_checker(self):
self._is_safety_checker_enabled = False
@property
def requires_safety_checker(self):
return getattr(self, "_is_safety_checker_enabled", self.config.enable_safety_checker)
def _encode_video(self, x):
return Cosmos3OmniPipeline._encode_video(self, x)
def decode_sound(self, latent):
return Cosmos3OmniPipeline.decode_sound(self, latent)
def _prepare_text_segment(self, input_ids, device):
return Cosmos3OmniPipeline._prepare_text_segment(self, input_ids, device)
def _prepare_vision_segment(self, *args, **kwargs):
return Cosmos3OmniPipeline._prepare_vision_segment(self, *args, **kwargs)
def _prepare_sound_segment(self, *args, **kwargs):
return Cosmos3OmniPipeline._prepare_sound_segment(self, *args, **kwargs)
def _prepare_action_segment(self, *args, **kwargs):
return Cosmos3OmniPipeline._prepare_action_segment(self, *args, **kwargs)
def _prepare_action_video_conditioning(self, *args, **kwargs):
return Cosmos3OmniPipeline._prepare_action_video_conditioning(self, *args, **kwargs)
def _remove_action_video_padding_from_latent(self, *args, **kwargs):
return Cosmos3OmniPipeline._remove_action_video_padding_from_latent(self, *args, **kwargs)
@staticmethod
def _build_action_json_prompt(*args, **kwargs):
return Cosmos3OmniPipeline._build_action_json_prompt(*args, **kwargs)
def tokenize_prompt(self, *args, **kwargs):
return Cosmos3OmniPipeline.tokenize_prompt(self, *args, **kwargs)
@staticmethod
def _mask_velocity_predictions(*args, **kwargs):
return Cosmos3OmniPipeline._mask_velocity_predictions(*args, **kwargs)
def _apply_video_safety_check(self, *args, **kwargs):
return Cosmos3OmniPipeline._apply_video_safety_check(self, *args, **kwargs)
class Cosmos3DistilledModularPipeline(Cosmos3OmniModularPipeline):
"""
A ModularPipeline for distilled (few-step) Cosmos 3 omni generation.
Distilled checkpoints bake classifier-free guidance into the weights and sample on a fixed schedule read from the
pipeline's `distilled_sigmas` config (populated from `modular_model_index.json`), so `guidance_scale` and
`num_inference_steps` are fixed and `negative_prompt` is not supported.
"""
default_blocks_name = "Cosmos3DistilledBlocks"
|