diff --git a/fastvideo/demo/gradio_web_demo.py b/fastvideo/demo/gradio_web_demo.py index 2c663c563..b294532d4 100644 --- a/fastvideo/demo/gradio_web_demo.py +++ b/fastvideo/demo/gradio_web_demo.py @@ -40,7 +40,7 @@ def load_model(args): mochi_genmo = True if mochi_genmo: - model_path = "/root/weights/dit.safetensors" + model_path = "/root/fastmochi_genmo/dit.safetensors" state_dcit = load_file(model_path) transformer = AsymmDiTJoint() transformer.load_state_dict(state_dcit) diff --git a/fastvideo/models/mochi_hf/modeling_mochi.py b/fastvideo/models/mochi_hf/modeling_mochi.py index e904fd8ba..9d5ff859c 100644 --- a/fastvideo/models/mochi_hf/modeling_mochi.py +++ b/fastvideo/models/mochi_hf/modeling_mochi.py @@ -55,11 +55,7 @@ from flash_attn import flash_attn_varlen_qkvpacked_func from flash_attn.bert_padding import pad_input, unpad_input from liger_kernel.ops.swiglu import LigerSiLUMulFunction -<<<<<<< HEAD from sageattention import sageattn -======= - ->>>>>>> de1e8d868e5b7e2a462f975a03a6043c01f2a8ba logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -188,14 +184,11 @@ class MochiAttention(nn.Module): attention_mask=attention_mask, **kwargs, ) -<<<<<<< HEAD def sage_attn_processor(query, key, value, attn_mask): query = query.squeeze(2) # B, S, H, D key = key.squeeze(2) value = value.squeeze(2) -======= ->>>>>>> de1e8d868e5b7e2a462f975a03a6043c01f2a8ba B = query.shape[0] assert B == 2, "This implementation expects batch size 2" @@ -330,17 +323,11 @@ class MochiAttnProcessor2_0: attn_mask = encoder_attention_mask[:, :].bool() attn_mask = F.pad(attn_mask, (sequence_length, 0), value=True) -<<<<<<< HEAD sage_attn = False if sage_attn: hidden_states = sage_attn_processor(query, key, value, attn_mask) else: hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) -======= - hidden_states = flash_attn_no_pad( - qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None - ) ->>>>>>> de1e8d868e5b7e2a462f975a03a6043c01f2a8ba # hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask = None, dropout_p=0.0, is_causal=False) diff --git a/fastvideo/sample/sample_t2v_mochi.py b/fastvideo/sample/sample_t2v_mochi.py index ac93fd51d..fad2dcb38 100644 --- a/fastvideo/sample/sample_t2v_mochi.py +++ b/fastvideo/sample/sample_t2v_mochi.py @@ -22,7 +22,8 @@ from typing import Dict from diffusers import FlowMatchEulerDiscreteScheduler from diffusers.utils import convert_unet_state_dict_to_peft from fastvideo.distill.solver import PCMFMScheduler - +from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.asymm_models_joint import AsymmDiTJoint +from safetensors.torch import load_file def initialize_distributed(): local_rank = int(os.getenv("RANK", 0)) @@ -53,12 +54,22 @@ def main(args): args.linear_threshold, args.linear_range, ) - if args.transformer_path is not None: - transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path) + + mochi_genmo = True + if mochi_genmo: + model_path = "/root/fastmochi_genmo/dit.safetensors" + state_dcit = load_file(model_path) + transformer = AsymmDiTJoint() + transformer.load_state_dict(state_dcit) + # from IPython import embed + # embed() + transformer.config.in_channels = 12 + print("load gennmo mochi successfully") else: - transformer = MochiTransformer3DModel.from_pretrained( - args.model_path, subfolder="transformer/" - ) + if args.transformer_path: + transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path) + else: + transformer = MochiTransformer3DModel.from_pretrained(args.model_path, subfolder='transformer/') pipe = MochiPipeline.from_pretrained( args.model_path, transformer=transformer, scheduler=scheduler