Spaces:
				
			
			
	
			
			
		Runtime error
		
	
	
	
			
			
	
	
	
	
		
		
		Runtime error
		
	Text to video files
Browse files- models/cldm_v15.yaml +81 -0
 - models/cldm_v15_no_cf_attn.yaml +81 -0
 - models/cldm_v21.yaml +85 -0
 - text_to_video/text_to_video_generator.py +77 -0
 - text_to_video/text_to_video_pipeline.py +550 -0
 
    	
        models/cldm_v15.yaml
    ADDED
    
    | 
         @@ -0,0 +1,81 @@ 
     | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
| 
         | 
|
| 1 | 
         
            +
            model:
         
     | 
| 2 | 
         
            +
              target: cldm.cldm.ControlLDM
         
     | 
| 3 | 
         
            +
              params:
         
     | 
| 4 | 
         
            +
                linear_start: 0.00085
         
     | 
| 5 | 
         
            +
                linear_end: 0.0120
         
     | 
| 6 | 
         
            +
                num_timesteps_cond: 1
         
     | 
| 7 | 
         
            +
                log_every_t: 200
         
     | 
| 8 | 
         
            +
                timesteps: 1000
         
     | 
| 9 | 
         
            +
                first_stage_key: "jpg"
         
     | 
| 10 | 
         
            +
                cond_stage_key: "txt"
         
     | 
| 11 | 
         
            +
                control_key: "hint"
         
     | 
| 12 | 
         
            +
                image_size: 64
         
     | 
| 13 | 
         
            +
                channels: 4
         
     | 
| 14 | 
         
            +
                cond_stage_trainable: false
         
     | 
| 15 | 
         
            +
                conditioning_key: crossattn
         
     | 
| 16 | 
         
            +
                monitor: val/loss_simple_ema
         
     | 
| 17 | 
         
            +
                scale_factor: 0.18215
         
     | 
| 18 | 
         
            +
                use_ema: False
         
     | 
| 19 | 
         
            +
                only_mid_control: False
         
     | 
| 20 | 
         
            +
             
     | 
| 21 | 
         
            +
                control_stage_config:
         
     | 
| 22 | 
         
            +
                  target: cldm.cldm.ControlNet
         
     | 
| 23 | 
         
            +
                  params:
         
     | 
| 24 | 
         
            +
                    image_size: 32 # unused
         
     | 
| 25 | 
         
            +
                    in_channels: 4
         
     | 
| 26 | 
         
            +
                    hint_channels: 3
         
     | 
| 27 | 
         
            +
                    model_channels: 320
         
     | 
| 28 | 
         
            +
                    attention_resolutions: [ 4, 2, 1 ]
         
     | 
| 29 | 
         
            +
                    num_res_blocks: 2
         
     | 
| 30 | 
         
            +
                    channel_mult: [ 1, 2, 4, 4 ]
         
     | 
| 31 | 
         
            +
                    num_heads: 8
         
     | 
| 32 | 
         
            +
                    use_spatial_transformer: True
         
     | 
| 33 | 
         
            +
                    transformer_depth: 1
         
     | 
| 34 | 
         
            +
                    context_dim: 768
         
     | 
| 35 | 
         
            +
                    use_checkpoint: True
         
     | 
| 36 | 
         
            +
                    use_cf_attn: True
         
     | 
| 37 | 
         
            +
                    legacy: False
         
     | 
| 38 | 
         
            +
             
     | 
| 39 | 
         
            +
                unet_config:
         
     | 
| 40 | 
         
            +
                  target: cldm.cldm.ControlledUnetModel
         
     | 
| 41 | 
         
            +
                  params:
         
     | 
| 42 | 
         
            +
                    image_size: 32 # unused
         
     | 
| 43 | 
         
            +
                    in_channels: 4
         
     | 
| 44 | 
         
            +
                    out_channels: 4
         
     | 
| 45 | 
         
            +
                    model_channels: 320
         
     | 
| 46 | 
         
            +
                    attention_resolutions: [ 4, 2, 1 ]
         
     | 
| 47 | 
         
            +
                    num_res_blocks: 2
         
     | 
| 48 | 
         
            +
                    channel_mult: [ 1, 2, 4, 4 ]
         
     | 
| 49 | 
         
            +
                    num_heads: 8
         
     | 
| 50 | 
         
            +
                    use_spatial_transformer: True
         
     | 
| 51 | 
         
            +
                    transformer_depth: 1
         
     | 
| 52 | 
         
            +
                    context_dim: 768
         
     | 
| 53 | 
         
            +
                    use_checkpoint: True
         
     | 
| 54 | 
         
            +
                    legacy: False
         
     | 
| 55 | 
         
            +
                    use_cf_attn: True
         
     | 
| 56 | 
         
            +
             
     | 
| 57 | 
         
            +
                first_stage_config:
         
     | 
| 58 | 
         
            +
                  target: ldm.models.autoencoder.AutoencoderKL
         
     | 
| 59 | 
         
            +
                  params:
         
     | 
| 60 | 
         
            +
                    embed_dim: 4
         
     | 
| 61 | 
         
            +
                    monitor: val/rec_loss
         
     | 
| 62 | 
         
            +
                    ddconfig:
         
     | 
| 63 | 
         
            +
                      double_z: true
         
     | 
| 64 | 
         
            +
                      z_channels: 4
         
     | 
| 65 | 
         
            +
                      resolution: 256
         
     | 
| 66 | 
         
            +
                      in_channels: 3
         
     | 
| 67 | 
         
            +
                      out_ch: 3
         
     | 
| 68 | 
         
            +
                      ch: 128
         
     | 
| 69 | 
         
            +
                      ch_mult:
         
     | 
| 70 | 
         
            +
                      - 1
         
     | 
| 71 | 
         
            +
                      - 2
         
     | 
| 72 | 
         
            +
                      - 4
         
     | 
| 73 | 
         
            +
                      - 4
         
     | 
| 74 | 
         
            +
                      num_res_blocks: 2
         
     | 
| 75 | 
         
            +
                      attn_resolutions: []
         
     | 
| 76 | 
         
            +
                      dropout: 0.0
         
     | 
| 77 | 
         
            +
                    lossconfig:
         
     | 
| 78 | 
         
            +
                      target: torch.nn.Identity
         
     | 
| 79 | 
         
            +
             
     | 
| 80 | 
         
            +
                cond_stage_config:
         
     | 
| 81 | 
         
            +
                  target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
         
     | 
    	
        models/cldm_v15_no_cf_attn.yaml
    ADDED
    
    | 
         @@ -0,0 +1,81 @@ 
     | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
| 
         | 
|
| 1 | 
         
            +
            model:
         
     | 
| 2 | 
         
            +
              target: cldm.cldm.ControlLDM
         
     | 
| 3 | 
         
            +
              params:
         
     | 
| 4 | 
         
            +
                linear_start: 0.00085
         
     | 
| 5 | 
         
            +
                linear_end: 0.0120
         
     | 
| 6 | 
         
            +
                num_timesteps_cond: 1
         
     | 
| 7 | 
         
            +
                log_every_t: 200
         
     | 
| 8 | 
         
            +
                timesteps: 1000
         
     | 
| 9 | 
         
            +
                first_stage_key: "jpg"
         
     | 
| 10 | 
         
            +
                cond_stage_key: "txt"
         
     | 
| 11 | 
         
            +
                control_key: "hint"
         
     | 
| 12 | 
         
            +
                image_size: 64
         
     | 
| 13 | 
         
            +
                channels: 4
         
     | 
| 14 | 
         
            +
                cond_stage_trainable: false
         
     | 
| 15 | 
         
            +
                conditioning_key: crossattn
         
     | 
| 16 | 
         
            +
                monitor: val/loss_simple_ema
         
     | 
| 17 | 
         
            +
                scale_factor: 0.18215
         
     | 
| 18 | 
         
            +
                use_ema: False
         
     | 
| 19 | 
         
            +
                only_mid_control: False
         
     | 
| 20 | 
         
            +
             
     | 
| 21 | 
         
            +
                control_stage_config:
         
     | 
| 22 | 
         
            +
                  target: cldm.cldm.ControlNet
         
     | 
| 23 | 
         
            +
                  params:
         
     | 
| 24 | 
         
            +
                    image_size: 32 # unused
         
     | 
| 25 | 
         
            +
                    in_channels: 4
         
     | 
| 26 | 
         
            +
                    hint_channels: 3
         
     | 
| 27 | 
         
            +
                    model_channels: 320
         
     | 
| 28 | 
         
            +
                    attention_resolutions: [ 4, 2, 1 ]
         
     | 
| 29 | 
         
            +
                    num_res_blocks: 2
         
     | 
| 30 | 
         
            +
                    channel_mult: [ 1, 2, 4, 4 ]
         
     | 
| 31 | 
         
            +
                    num_heads: 8
         
     | 
| 32 | 
         
            +
                    use_spatial_transformer: True
         
     | 
| 33 | 
         
            +
                    transformer_depth: 1
         
     | 
| 34 | 
         
            +
                    context_dim: 768
         
     | 
| 35 | 
         
            +
                    use_checkpoint: True
         
     | 
| 36 | 
         
            +
                    use_cf_attn: False
         
     | 
| 37 | 
         
            +
                    legacy: False
         
     | 
| 38 | 
         
            +
             
     | 
| 39 | 
         
            +
                unet_config:
         
     | 
| 40 | 
         
            +
                  target: cldm.cldm.ControlledUnetModel
         
     | 
| 41 | 
         
            +
                  params:
         
     | 
| 42 | 
         
            +
                    image_size: 32 # unused
         
     | 
| 43 | 
         
            +
                    in_channels: 4
         
     | 
| 44 | 
         
            +
                    out_channels: 4
         
     | 
| 45 | 
         
            +
                    model_channels: 320
         
     | 
| 46 | 
         
            +
                    attention_resolutions: [ 4, 2, 1 ]
         
     | 
| 47 | 
         
            +
                    num_res_blocks: 2
         
     | 
| 48 | 
         
            +
                    channel_mult: [ 1, 2, 4, 4 ]
         
     | 
| 49 | 
         
            +
                    num_heads: 8
         
     | 
| 50 | 
         
            +
                    use_spatial_transformer: True
         
     | 
| 51 | 
         
            +
                    transformer_depth: 1
         
     | 
| 52 | 
         
            +
                    context_dim: 768
         
     | 
| 53 | 
         
            +
                    use_checkpoint: True
         
     | 
| 54 | 
         
            +
                    legacy: False
         
     | 
| 55 | 
         
            +
                    use_cf_attn: False
         
     | 
| 56 | 
         
            +
             
     | 
| 57 | 
         
            +
                first_stage_config:
         
     | 
| 58 | 
         
            +
                  target: ldm.models.autoencoder.AutoencoderKL
         
     | 
| 59 | 
         
            +
                  params:
         
     | 
| 60 | 
         
            +
                    embed_dim: 4
         
     | 
| 61 | 
         
            +
                    monitor: val/rec_loss
         
     | 
| 62 | 
         
            +
                    ddconfig:
         
     | 
| 63 | 
         
            +
                      double_z: true
         
     | 
| 64 | 
         
            +
                      z_channels: 4
         
     | 
| 65 | 
         
            +
                      resolution: 256
         
     | 
| 66 | 
         
            +
                      in_channels: 3
         
     | 
| 67 | 
         
            +
                      out_ch: 3
         
     | 
| 68 | 
         
            +
                      ch: 128
         
     | 
| 69 | 
         
            +
                      ch_mult:
         
     | 
| 70 | 
         
            +
                      - 1
         
     | 
| 71 | 
         
            +
                      - 2
         
     | 
| 72 | 
         
            +
                      - 4
         
     | 
| 73 | 
         
            +
                      - 4
         
     | 
| 74 | 
         
            +
                      num_res_blocks: 2
         
     | 
| 75 | 
         
            +
                      attn_resolutions: []
         
     | 
| 76 | 
         
            +
                      dropout: 0.0
         
     | 
| 77 | 
         
            +
                    lossconfig:
         
     | 
| 78 | 
         
            +
                      target: torch.nn.Identity
         
     | 
| 79 | 
         
            +
             
     | 
| 80 | 
         
            +
                cond_stage_config:
         
     | 
| 81 | 
         
            +
                  target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
         
     | 
    	
        models/cldm_v21.yaml
    ADDED
    
    | 
         @@ -0,0 +1,85 @@ 
     | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
| 
         | 
|
| 1 | 
         
            +
            model:
         
     | 
| 2 | 
         
            +
              target: cldm.cldm.ControlLDM
         
     | 
| 3 | 
         
            +
              params:
         
     | 
| 4 | 
         
            +
                linear_start: 0.00085
         
     | 
| 5 | 
         
            +
                linear_end: 0.0120
         
     | 
| 6 | 
         
            +
                num_timesteps_cond: 1
         
     | 
| 7 | 
         
            +
                log_every_t: 200
         
     | 
| 8 | 
         
            +
                timesteps: 1000
         
     | 
| 9 | 
         
            +
                first_stage_key: "jpg"
         
     | 
| 10 | 
         
            +
                cond_stage_key: "txt"
         
     | 
| 11 | 
         
            +
                control_key: "hint"
         
     | 
| 12 | 
         
            +
                image_size: 64
         
     | 
| 13 | 
         
            +
                channels: 4
         
     | 
| 14 | 
         
            +
                cond_stage_trainable: false
         
     | 
| 15 | 
         
            +
                conditioning_key: crossattn
         
     | 
| 16 | 
         
            +
                monitor: val/loss_simple_ema
         
     | 
| 17 | 
         
            +
                scale_factor: 0.18215
         
     | 
| 18 | 
         
            +
                use_ema: False
         
     | 
| 19 | 
         
            +
                only_mid_control: False
         
     | 
| 20 | 
         
            +
             
     | 
| 21 | 
         
            +
                control_stage_config:
         
     | 
| 22 | 
         
            +
                  target: cldm.cldm.ControlNet
         
     | 
| 23 | 
         
            +
                  params:
         
     | 
| 24 | 
         
            +
                    use_checkpoint: True
         
     | 
| 25 | 
         
            +
                    image_size: 32 # unused
         
     | 
| 26 | 
         
            +
                    in_channels: 4
         
     | 
| 27 | 
         
            +
                    hint_channels: 3
         
     | 
| 28 | 
         
            +
                    model_channels: 320
         
     | 
| 29 | 
         
            +
                    attention_resolutions: [ 4, 2, 1 ]
         
     | 
| 30 | 
         
            +
                    num_res_blocks: 2
         
     | 
| 31 | 
         
            +
                    channel_mult: [ 1, 2, 4, 4 ]
         
     | 
| 32 | 
         
            +
                    num_head_channels: 64 # need to fix for flash-attn
         
     | 
| 33 | 
         
            +
                    use_spatial_transformer: True
         
     | 
| 34 | 
         
            +
                    use_linear_in_transformer: True
         
     | 
| 35 | 
         
            +
                    transformer_depth: 1
         
     | 
| 36 | 
         
            +
                    context_dim: 1024
         
     | 
| 37 | 
         
            +
                    legacy: False
         
     | 
| 38 | 
         
            +
             
     | 
| 39 | 
         
            +
                unet_config:
         
     | 
| 40 | 
         
            +
                  target: cldm.cldm.ControlledUnetModel
         
     | 
| 41 | 
         
            +
                  params:
         
     | 
| 42 | 
         
            +
                    use_checkpoint: True
         
     | 
| 43 | 
         
            +
                    image_size: 32 # unused
         
     | 
| 44 | 
         
            +
                    in_channels: 4
         
     | 
| 45 | 
         
            +
                    out_channels: 4
         
     | 
| 46 | 
         
            +
                    model_channels: 320
         
     | 
| 47 | 
         
            +
                    attention_resolutions: [ 4, 2, 1 ]
         
     | 
| 48 | 
         
            +
                    num_res_blocks: 2
         
     | 
| 49 | 
         
            +
                    channel_mult: [ 1, 2, 4, 4 ]
         
     | 
| 50 | 
         
            +
                    num_head_channels: 64 # need to fix for flash-attn
         
     | 
| 51 | 
         
            +
                    use_spatial_transformer: True
         
     | 
| 52 | 
         
            +
                    use_linear_in_transformer: True
         
     | 
| 53 | 
         
            +
                    transformer_depth: 1
         
     | 
| 54 | 
         
            +
                    context_dim: 1024
         
     | 
| 55 | 
         
            +
                    legacy: False
         
     | 
| 56 | 
         
            +
             
     | 
| 57 | 
         
            +
                first_stage_config:
         
     | 
| 58 | 
         
            +
                  target: ldm.models.autoencoder.AutoencoderKL
         
     | 
| 59 | 
         
            +
                  params:
         
     | 
| 60 | 
         
            +
                    embed_dim: 4
         
     | 
| 61 | 
         
            +
                    monitor: val/rec_loss
         
     | 
| 62 | 
         
            +
                    ddconfig:
         
     | 
| 63 | 
         
            +
                      #attn_type: "vanilla-xformers"
         
     | 
| 64 | 
         
            +
                      double_z: true
         
     | 
| 65 | 
         
            +
                      z_channels: 4
         
     | 
| 66 | 
         
            +
                      resolution: 256
         
     | 
| 67 | 
         
            +
                      in_channels: 3
         
     | 
| 68 | 
         
            +
                      out_ch: 3
         
     | 
| 69 | 
         
            +
                      ch: 128
         
     | 
| 70 | 
         
            +
                      ch_mult:
         
     | 
| 71 | 
         
            +
                      - 1
         
     | 
| 72 | 
         
            +
                      - 2
         
     | 
| 73 | 
         
            +
                      - 4
         
     | 
| 74 | 
         
            +
                      - 4
         
     | 
| 75 | 
         
            +
                      num_res_blocks: 2
         
     | 
| 76 | 
         
            +
                      attn_resolutions: []
         
     | 
| 77 | 
         
            +
                      dropout: 0.0
         
     | 
| 78 | 
         
            +
                    lossconfig:
         
     | 
| 79 | 
         
            +
                      target: torch.nn.Identity
         
     | 
| 80 | 
         
            +
             
     | 
| 81 | 
         
            +
                cond_stage_config:
         
     | 
| 82 | 
         
            +
                  target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder
         
     | 
| 83 | 
         
            +
                  params:
         
     | 
| 84 | 
         
            +
                    freeze: True
         
     | 
| 85 | 
         
            +
                    layer: "penultimate"
         
     | 
    	
        text_to_video/text_to_video_generator.py
    ADDED
    
    | 
         @@ -0,0 +1,77 @@ 
     | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
| 
         | 
|
| 1 | 
         
            +
            from text_to_video.tuneavideo.pipelines.pipeline_text_to_video import TuneAVideoPipeline
         
     | 
| 2 | 
         
            +
            from text_to_video.tuneavideo.models.unet import UNet3DConditionModel
         
     | 
| 3 | 
         
            +
            import torch
         
     | 
| 4 | 
         
            +
            from diffusers import AutoencoderKL, DDIMScheduler
         
     | 
| 5 | 
         
            +
            from transformers import CLIPTextModel, CLIPTokenizer
         
     | 
| 6 | 
         
            +
             
     | 
| 7 | 
         
            +
             
     | 
| 8 | 
         
            +
            class TextToVideo():
         
     | 
| 9 | 
         
            +
             
     | 
| 10 | 
         
            +
                
         
     | 
| 11 | 
         
            +
                def __init__(self,sd_path = None,motion_field_strength = 12, video_length = 8,t0 = 881, t1=941,use_cf_attn=True,use_motion_field=True) -> None:
         
     | 
| 12 | 
         
            +
                    g = torch.Generator(device='cuda')
         
     | 
| 13 | 
         
            +
                    g.manual_seed(22)
         
     | 
| 14 | 
         
            +
                    self.g = g
         
     | 
| 15 | 
         
            +
             
     | 
| 16 | 
         
            +
                    print(f"Loading model SD-Net model file from {sd_path}")
         
     | 
| 17 | 
         
            +
             
     | 
| 18 | 
         
            +
                    self.dtype = torch.float16
         
     | 
| 19 | 
         
            +
                    noise_scheduler = DDIMScheduler.from_pretrained(
         
     | 
| 20 | 
         
            +
                        sd_path, subfolder="scheduler")
         
     | 
| 21 | 
         
            +
                    tokenizer = CLIPTokenizer.from_pretrained(
         
     | 
| 22 | 
         
            +
                        sd_path, subfolder="tokenizer")
         
     | 
| 23 | 
         
            +
                    text_encoder = CLIPTextModel.from_pretrained(
         
     | 
| 24 | 
         
            +
                        sd_path, subfolder="text_encoder")
         
     | 
| 25 | 
         
            +
                    vae = AutoencoderKL.from_pretrained(sd_path, subfolder="vae")
         
     | 
| 26 | 
         
            +
             
     | 
| 27 | 
         
            +
             
     | 
| 28 | 
         
            +
                    unet = UNet3DConditionModel.from_pretrained_2d(
         
     | 
| 29 | 
         
            +
                        sd_path, subfolder="unet", use_cf_attn=use_cf_attn)
         
     | 
| 30 | 
         
            +
                    self.pipe = TuneAVideoPipeline(
         
     | 
| 31 | 
         
            +
                        vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet,
         
     | 
| 32 | 
         
            +
                        scheduler=DDIMScheduler.from_pretrained(
         
     | 
| 33 | 
         
            +
                            sd_path, subfolder="scheduler")
         
     | 
| 34 | 
         
            +
                    ).to('cuda').to(self.dtype)
         
     | 
| 35 | 
         
            +
             
     | 
| 36 | 
         
            +
                    noise_scheduler.set_timesteps(50, device='cuda')
         
     | 
| 37 | 
         
            +
             
     | 
| 38 | 
         
            +
                    # t0 parameter (DDIM backward from noise until t0)
         
     | 
| 39 | 
         
            +
                    self.t0 = t0
         
     | 
| 40 | 
         
            +
                    
         
     | 
| 41 | 
         
            +
             
     | 
| 42 | 
         
            +
                    # from t0 apply DDPM forward until t1
         
     | 
| 43 | 
         
            +
                    self.t1 = t1
         
     | 
| 44 | 
         
            +
             
     | 
| 45 | 
         
            +
                    self.use_foreground_motion_field = False  # apply motion field on forground object (not used)
         
     | 
| 46 | 
         
            +
             
     | 
| 47 | 
         
            +
                    # strength of motion field (delta_x = delta_y in Sect 3.3.1)
         
     | 
| 48 | 
         
            +
                    self.motion_field_strength = motion_field_strength
         
     | 
| 49 | 
         
            +
                    self.use_motion_field = use_motion_field  # apply general motion field
         
     | 
| 50 | 
         
            +
                    self.smooth_bg = False  # temporally smooth background
         
     | 
| 51 | 
         
            +
                    self.smooth_bg_strength = 0.4  # alpha = (1-self.smooth_bg_strength) in Eq (9)
         
     | 
| 52 | 
         
            +
             
     | 
| 53 | 
         
            +
             
     | 
| 54 | 
         
            +
                    self.video_length = video_length
         
     | 
| 55 | 
         
            +
                    
         
     | 
| 56 | 
         
            +
                def inference(self, prompt):
         
     | 
| 57 | 
         
            +
             
     | 
| 58 | 
         
            +
                    prompt_compute = [prompt]
         
     | 
| 59 | 
         
            +
                    xT = torch.randn((1, 4, 1, 64, 64), dtype=self.dtype, device="cuda")
         
     | 
| 60 | 
         
            +
                    result = self.pipe(prompt_compute,
         
     | 
| 61 | 
         
            +
                                       video_length=self.video_length,
         
     | 
| 62 | 
         
            +
                                       height=512,
         
     | 
| 63 | 
         
            +
                                       width=512,
         
     | 
| 64 | 
         
            +
                                       num_inference_steps=50,
         
     | 
| 65 | 
         
            +
                                       guidance_scale=7.5,
         
     | 
| 66 | 
         
            +
                                       guidance_stop_step=1.0,
         
     | 
| 67 | 
         
            +
                                       t0=self.t0,
         
     | 
| 68 | 
         
            +
                                       t1=self.t1,
         
     | 
| 69 | 
         
            +
                                       xT=xT,
         
     | 
| 70 | 
         
            +
                                       use_foreground_motion_field=self.use_foreground_motion_field,
         
     | 
| 71 | 
         
            +
                                       motion_field_strength=self.motion_field_strength,
         
     | 
| 72 | 
         
            +
                                       use_motion_field=self.use_motion_field,
         
     | 
| 73 | 
         
            +
                                       smooth_bg=self.smooth_bg,
         
     | 
| 74 | 
         
            +
                                       smooth_bg_strength=self.smooth_bg_strength,
         
     | 
| 75 | 
         
            +
                                       generator=self.g)
         
     | 
| 76 | 
         
            +
             
     | 
| 77 | 
         
            +
                    return result.videos[0]
         
     | 
    	
        text_to_video/text_to_video_pipeline.py
    ADDED
    
    | 
         @@ -0,0 +1,550 @@ 
     | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
| 
         | 
|
| 1 | 
         
            +
            from diffusers import StableDiffusionPipeline
         
     | 
| 2 | 
         
            +
            import torch
         
     | 
| 3 | 
         
            +
            from dataclasses import dataclass
         
     | 
| 4 | 
         
            +
            from typing import Callable, List, Optional, Union
         
     | 
| 5 | 
         
            +
            import numpy as np
         
     | 
| 6 | 
         
            +
            from diffusers.utils import deprecate, logging, BaseOutput
         
     | 
| 7 | 
         
            +
            from einops import rearrange, repeat
         
     | 
| 8 | 
         
            +
            from torch.nn.functional import grid_sample
         
     | 
| 9 | 
         
            +
            import torchvision.transforms as T
         
     | 
| 10 | 
         
            +
            from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
         
     | 
| 11 | 
         
            +
            from diffusers.models import AutoencoderKL, UNet2DConditionModel
         
     | 
| 12 | 
         
            +
            from diffusers.schedulers import KarrasDiffusionSchedulers
         
     | 
| 13 | 
         
            +
            from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
         
     | 
| 14 | 
         
            +
             
     | 
| 15 | 
         
            +
            @dataclass
         
     | 
| 16 | 
         
            +
            class TextToVideoPipelineOutput(BaseOutput):
         
     | 
| 17 | 
         
            +
                videos: Union[torch.Tensor, np.ndarray]
         
     | 
| 18 | 
         
            +
                code: Union[torch.Tensor, np.ndarray]
         
     | 
| 19 | 
         
            +
             
     | 
| 20 | 
         
            +
             
     | 
| 21 | 
         
            +
             
     | 
| 22 | 
         
            +
            def coords_grid(batch, ht, wd, device):
         
     | 
| 23 | 
         
            +
                # Adapted from https://github.com/princeton-vl/RAFT/blob/master/core/utils/utils.py
         
     | 
| 24 | 
         
            +
                coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device))
         
     | 
| 25 | 
         
            +
                coords = torch.stack(coords[::-1], dim=0).float()
         
     | 
| 26 | 
         
            +
                return coords[None].repeat(batch, 1, 1, 1)
         
     | 
| 27 | 
         
            +
             
     | 
| 28 | 
         
            +
             
     | 
| 29 | 
         
            +
             
     | 
| 30 | 
         
            +
            class TextToVideoPipeline(StableDiffusionPipeline):
         
     | 
| 31 | 
         
            +
                def __init__(
         
     | 
| 32 | 
         
            +
                    self,
         
     | 
| 33 | 
         
            +
                    vae: AutoencoderKL,
         
     | 
| 34 | 
         
            +
                    text_encoder: CLIPTextModel,
         
     | 
| 35 | 
         
            +
                    tokenizer: CLIPTokenizer,
         
     | 
| 36 | 
         
            +
                    unet: UNet2DConditionModel,
         
     | 
| 37 | 
         
            +
                    scheduler: KarrasDiffusionSchedulers,
         
     | 
| 38 | 
         
            +
                    safety_checker: StableDiffusionSafetyChecker,
         
     | 
| 39 | 
         
            +
                    feature_extractor: CLIPFeatureExtractor,
         
     | 
| 40 | 
         
            +
                    requires_safety_checker: bool = True,
         
     | 
| 41 | 
         
            +
                    ):
         
     | 
| 42 | 
         
            +
                    #super().__init__(*args,**kwargs)
         
     | 
| 43 | 
         
            +
                    super().__init__(vae,text_encoder,tokenizer,unet,scheduler,safety_checker,feature_extractor,requires_safety_checker)
         
     | 
| 44 | 
         
            +
             
     | 
| 45 | 
         
            +
             
     | 
| 46 | 
         
            +
                def DDPM_forward(self, x0, t0, tMax, generator, device, shape, text_embeddings):
         
     | 
| 47 | 
         
            +
                    rand_device = "cpu" if device.type == "mps" else device
         
     | 
| 48 | 
         
            +
               
         
     | 
| 49 | 
         
            +
                    if x0 is None:
         
     | 
| 50 | 
         
            +
                        return torch.randn(shape, generator=generator, device=rand_device, dtype=text_embeddings.dtype).to(device)
         
     | 
| 51 | 
         
            +
                    else:
         
     | 
| 52 | 
         
            +
                        eps = torch.randn_like(x0, dtype=text_embeddings.dtype).to(device)
         
     | 
| 53 | 
         
            +
                        alpha_vec = torch.prod(self.scheduler.alphas[t0:tMax])
         
     | 
| 54 | 
         
            +
                        xt = torch.sqrt(alpha_vec) * x0 + \
         
     | 
| 55 | 
         
            +
                            torch.sqrt(1-alpha_vec) * eps
         
     | 
| 56 | 
         
            +
                        return xt
         
     | 
| 57 | 
         
            +
             
     | 
| 58 | 
         
            +
             
     | 
| 59 | 
         
            +
                def prepare_latents(self, batch_size, num_channels_latents, video_length, height, width, dtype, device, generator, latents=None):
         
     | 
| 60 | 
         
            +
                    shape = (batch_size, num_channels_latents, video_length, height //
         
     | 
| 61 | 
         
            +
                             self.vae_scale_factor, width // self.vae_scale_factor)
         
     | 
| 62 | 
         
            +
                    if isinstance(generator, list) and len(generator) != batch_size:
         
     | 
| 63 | 
         
            +
                        raise ValueError(
         
     | 
| 64 | 
         
            +
                            f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
         
     | 
| 65 | 
         
            +
                            f" size of {batch_size}. Make sure the batch size matches the length of the generators."
         
     | 
| 66 | 
         
            +
                        )
         
     | 
| 67 | 
         
            +
             
     | 
| 68 | 
         
            +
                    if latents is None:
         
     | 
| 69 | 
         
            +
                        rand_device = "cpu" if device.type == "mps" else device
         
     | 
| 70 | 
         
            +
             
     | 
| 71 | 
         
            +
                        if isinstance(generator, list):
         
     | 
| 72 | 
         
            +
                            shape = (1,) + shape[1:]
         
     | 
| 73 | 
         
            +
                            latents = [
         
     | 
| 74 | 
         
            +
                                torch.randn(
         
     | 
| 75 | 
         
            +
                                    shape, generator=generator[i], device=rand_device, dtype=dtype)
         
     | 
| 76 | 
         
            +
                                for i in range(batch_size)
         
     | 
| 77 | 
         
            +
                            ]
         
     | 
| 78 | 
         
            +
                            latents = torch.cat(latents, dim=0).to(device)
         
     | 
| 79 | 
         
            +
                        else:
         
     | 
| 80 | 
         
            +
                            latents = torch.randn(
         
     | 
| 81 | 
         
            +
                                shape, generator=generator, device=rand_device, dtype=dtype).to(device)
         
     | 
| 82 | 
         
            +
                    else:
         
     | 
| 83 | 
         
            +
                        latents = latents.to(device)
         
     | 
| 84 | 
         
            +
             
     | 
| 85 | 
         
            +
                    # scale the initial noise by the standard deviation required by the scheduler
         
     | 
| 86 | 
         
            +
                    latents = latents * self.scheduler.init_noise_sigma
         
     | 
| 87 | 
         
            +
                    return latents
         
     | 
| 88 | 
         
            +
             
     | 
| 89 | 
         
            +
             
     | 
| 90 | 
         
            +
             
     | 
| 91 | 
         
            +
                def warp_latents(self, latents, reference_flow):
         
     | 
| 92 | 
         
            +
                    _, _, H, W = reference_flow.size()
         
     | 
| 93 | 
         
            +
                    b, c, f, h, w = latents.size()
         
     | 
| 94 | 
         
            +
                    coords0 = coords_grid(f, H, W, device=latents.device).to(latents.dtype)
         
     | 
| 95 | 
         
            +
                    coords_t0 = coords0 + reference_flow
         
     | 
| 96 | 
         
            +
                    coords_t0[:, 0] /= W
         
     | 
| 97 | 
         
            +
                    coords_t0[:, 1] /= H
         
     | 
| 98 | 
         
            +
                    coords_t0 = coords_t0 * 2.0 - 1.0
         
     | 
| 99 | 
         
            +
                    coords_t0 = T.Resize((h, w))(coords_t0)
         
     | 
| 100 | 
         
            +
                    coords_t0 = rearrange(coords_t0, 'f c h w -> f h w c')
         
     | 
| 101 | 
         
            +
                    latents_0 = latents[:, :, 0]
         
     | 
| 102 | 
         
            +
                    latents_0 = latents_0.repeat(f, 1, 1, 1)
         
     | 
| 103 | 
         
            +
                    warped = grid_sample(latents_0, coords_t0,
         
     | 
| 104 | 
         
            +
                                            mode='nearest', padding_mode='reflection')
         
     | 
| 105 | 
         
            +
                    warped = rearrange(warped, '(b f) c h w -> b c f h w', f=f)
         
     | 
| 106 | 
         
            +
                    return warped
         
     | 
| 107 | 
         
            +
             
     | 
| 108 | 
         
            +
                def warp_latents_independently(self, latents, reference_flow):
         
     | 
| 109 | 
         
            +
                    _, _, H, W = reference_flow.size()
         
     | 
| 110 | 
         
            +
                    b, c, f, h, w = latents.size()
         
     | 
| 111 | 
         
            +
                    assert b == 1
         
     | 
| 112 | 
         
            +
                    coords0 = coords_grid(f, H, W, device=latents.device).to(latents.dtype)
         
     | 
| 113 | 
         
            +
                    coords_t0 = coords0 + reference_flow
         
     | 
| 114 | 
         
            +
             
     | 
| 115 | 
         
            +
                    coords_t0[:, 0] /= W
         
     | 
| 116 | 
         
            +
                    coords_t0[:, 1] /= H
         
     | 
| 117 | 
         
            +
                    coords_t0 = coords_t0 * 2.0 - 1.0
         
     | 
| 118 | 
         
            +
             
     | 
| 119 | 
         
            +
                    coords_t0 = T.Resize((h, w))(coords_t0)
         
     | 
| 120 | 
         
            +
             
     | 
| 121 | 
         
            +
                    coords_t0 = rearrange(coords_t0, 'f c h w -> f h w c')
         
     | 
| 122 | 
         
            +
             
     | 
| 123 | 
         
            +
                    latents_0 = rearrange(latents[0], 'c f h w -> f  c  h w')
         
     | 
| 124 | 
         
            +
             
     | 
| 125 | 
         
            +
                    warped = grid_sample(latents_0, coords_t0,
         
     | 
| 126 | 
         
            +
                                         mode='nearest', padding_mode='reflection')
         
     | 
| 127 | 
         
            +
                    warped = rearrange(warped, '(b f) c h w -> b c f h w', f=f)
         
     | 
| 128 | 
         
            +
                    return warped
         
     | 
| 129 | 
         
            +
             
     | 
| 130 | 
         
            +
                def DDIM_backward(self, num_inference_steps, timesteps, skip_t, t0, t1, do_classifier_free_guidance, null_embs, text_embeddings, latents_local, latents_dtype, guidance_scale, guidance_stop_step, callback, callback_steps, extra_step_kwargs, num_warmup_steps):
         
     | 
| 131 | 
         
            +
                    entered = False
         
     | 
| 132 | 
         
            +
                    
         
     | 
| 133 | 
         
            +
                    f = latents_local.shape[2]
         
     | 
| 134 | 
         
            +
                    latents_local = rearrange(latents_local,"b c f w h -> (b f) c w h")
         
     | 
| 135 | 
         
            +
                    
         
     | 
| 136 | 
         
            +
                    latents = latents_local.detach().clone()
         
     | 
| 137 | 
         
            +
                    x_t0_1 = None
         
     | 
| 138 | 
         
            +
                    x_t1_1 = None
         
     | 
| 139 | 
         
            +
                    
         
     | 
| 140 | 
         
            +
                    
         
     | 
| 141 | 
         
            +
             
     | 
| 142 | 
         
            +
                    with self.progress_bar(total=num_inference_steps) as progress_bar:
         
     | 
| 143 | 
         
            +
                        for i, t in enumerate(timesteps):
         
     | 
| 144 | 
         
            +
                            if t > skip_t:
         
     | 
| 145 | 
         
            +
                                # print("Skipping frame!")
         
     | 
| 146 | 
         
            +
                                continue
         
     | 
| 147 | 
         
            +
                            else:
         
     | 
| 148 | 
         
            +
                                if not entered:
         
     | 
| 149 | 
         
            +
                                    print(
         
     | 
| 150 | 
         
            +
                                        f"Continue DDIM with i = {i}, t = {t}, latent = {latents.shape}, device = {latents.device}, type = {latents.dtype}")
         
     | 
| 151 | 
         
            +
                                    entered = True
         
     | 
| 152 | 
         
            +
             
     | 
| 153 | 
         
            +
                            latents = latents.detach()
         
     | 
| 154 | 
         
            +
                            # expand the latents if we are doing classifier free guidance
         
     | 
| 155 | 
         
            +
                            latent_model_input = torch.cat(
         
     | 
| 156 | 
         
            +
                                [latents] * 2) if do_classifier_free_guidance else latents
         
     | 
| 157 | 
         
            +
                            latent_model_input = self.scheduler.scale_model_input(
         
     | 
| 158 | 
         
            +
                                latent_model_input, t)
         
     | 
| 159 | 
         
            +
             
     | 
| 160 | 
         
            +
                            # predict the noise residual
         
     | 
| 161 | 
         
            +
                            with torch.no_grad():
         
     | 
| 162 | 
         
            +
                                if null_embs is not None:
         
     | 
| 163 | 
         
            +
                                    text_embeddings[0] = null_embs[i][0]
         
     | 
| 164 | 
         
            +
                                te = torch.cat([repeat(text_embeddings[0,:,:], "c k -> f c k",f=f),repeat(text_embeddings[1,:,:], "c k -> f c k",f=f)]) 
         
     | 
| 165 | 
         
            +
                                noise_pred = self.unet(
         
     | 
| 166 | 
         
            +
                                    latent_model_input, t, encoder_hidden_states=te).sample.to(dtype=latents_dtype)
         
     | 
| 167 | 
         
            +
             
     | 
| 168 | 
         
            +
                            # perform guidance
         
     | 
| 169 | 
         
            +
                            if do_classifier_free_guidance:
         
     | 
| 170 | 
         
            +
                                noise_pred_uncond, noise_pred_text = noise_pred.chunk(
         
     | 
| 171 | 
         
            +
                                    2)
         
     | 
| 172 | 
         
            +
                                noise_pred = noise_pred_uncond + guidance_scale * \
         
     | 
| 173 | 
         
            +
                                    (noise_pred_text - noise_pred_uncond)
         
     | 
| 174 | 
         
            +
             
     | 
| 175 | 
         
            +
                            if i >= guidance_stop_step * len(timesteps):
         
     | 
| 176 | 
         
            +
                                alpha = 0
         
     | 
| 177 | 
         
            +
                            # compute the previous noisy sample x_t -> x_t-1
         
     | 
| 178 | 
         
            +
                            latents = self.scheduler.step(
         
     | 
| 179 | 
         
            +
                                noise_pred, t, latents, **extra_step_kwargs).prev_sample
         
     | 
| 180 | 
         
            +
                            # latents = latents - alpha * grads / (torch.norm(grads) + 1e-10)
         
     | 
| 181 | 
         
            +
                            # call the callback, if provided
         
     | 
| 182 | 
         
            +
             
     | 
| 183 | 
         
            +
                            if i < len(timesteps)-1 and timesteps[i+1] == t0:
         
     | 
| 184 | 
         
            +
                                x_t0_1 = latents.detach().clone()
         
     | 
| 185 | 
         
            +
                                print(f"latent t0 found at i = {i}, t = {t}")
         
     | 
| 186 | 
         
            +
                            elif i < len(timesteps)-1 and timesteps[i+1] == t1:
         
     | 
| 187 | 
         
            +
                                x_t1_1 = latents.detach().clone()
         
     | 
| 188 | 
         
            +
                                print(f"latent t1 found at i={i}, t = {t}")
         
     | 
| 189 | 
         
            +
             
     | 
| 190 | 
         
            +
                            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
         
     | 
| 191 | 
         
            +
                                progress_bar.update()
         
     | 
| 192 | 
         
            +
                                if callback is not None and i % callback_steps == 0:
         
     | 
| 193 | 
         
            +
                                    callback(i, t, latents)
         
     | 
| 194 | 
         
            +
             
     | 
| 195 | 
         
            +
                    
         
     | 
| 196 | 
         
            +
                    latents = rearrange(latents,"(b f) c w h -> b c f  w h",f = f)
         
     | 
| 197 | 
         
            +
                    
         
     | 
| 198 | 
         
            +
                   
         
     | 
| 199 | 
         
            +
                    
         
     | 
| 200 | 
         
            +
                    res = {"x0": latents.detach().clone()}
         
     | 
| 201 | 
         
            +
                    if x_t0_1 is not None:
         
     | 
| 202 | 
         
            +
                        x_t0_1 = rearrange(x_t0_1,"(b f) c w h -> b c f  w h",f = f)
         
     | 
| 203 | 
         
            +
                        res["x_t0_1"] = x_t0_1.detach().clone()
         
     | 
| 204 | 
         
            +
                    if x_t1_1 is not None:
         
     | 
| 205 | 
         
            +
                        x_t1_1 = rearrange(x_t1_1,"(b f) c w h -> b c f  w h",f = f)
         
     | 
| 206 | 
         
            +
                        res["x_t1_1"] = x_t1_1.detach().clone()
         
     | 
| 207 | 
         
            +
                    return res
         
     | 
| 208 | 
         
            +
             
     | 
| 209 | 
         
            +
                def decode_latents(self, latents):
         
     | 
| 210 | 
         
            +
                    video_length = latents.shape[2]
         
     | 
| 211 | 
         
            +
                    latents = 1 / 0.18215 * latents
         
     | 
| 212 | 
         
            +
                    latents = rearrange(latents, "b c f h w -> (b f) c h w")
         
     | 
| 213 | 
         
            +
                    video = self.vae.decode(latents).sample
         
     | 
| 214 | 
         
            +
                    video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length)
         
     | 
| 215 | 
         
            +
                    video = (video / 2 + 0.5).clamp(0, 1)
         
     | 
| 216 | 
         
            +
                    # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
         
     | 
| 217 | 
         
            +
                    return video
         
     | 
| 218 | 
         
            +
             
     | 
| 219 | 
         
            +
             
     | 
| 220 | 
         
            +
             
     | 
| 221 | 
         
            +
                @torch.no_grad()
         
     | 
| 222 | 
         
            +
                def __call__(
         
     | 
| 223 | 
         
            +
                    self,
         
     | 
| 224 | 
         
            +
                    prompt: Union[str, List[str]],
         
     | 
| 225 | 
         
            +
                    video_length: Optional[int],
         
     | 
| 226 | 
         
            +
                    height: Optional[int] = None,
         
     | 
| 227 | 
         
            +
                    width: Optional[int] = None,
         
     | 
| 228 | 
         
            +
                    num_inference_steps: int = 50,
         
     | 
| 229 | 
         
            +
                    guidance_scale: float = 7.5,
         
     | 
| 230 | 
         
            +
                    guidance_stop_step: float = 0.5,
         
     | 
| 231 | 
         
            +
                    negative_prompt: Optional[Union[str, List[str]]] = None,
         
     | 
| 232 | 
         
            +
                    num_videos_per_prompt: Optional[int] = 1,
         
     | 
| 233 | 
         
            +
                    eta: float = 0.0,
         
     | 
| 234 | 
         
            +
                    generator: Optional[Union[torch.Generator,
         
     | 
| 235 | 
         
            +
                                              List[torch.Generator]]] = None,
         
     | 
| 236 | 
         
            +
                    xT: Optional[torch.FloatTensor] = None,
         
     | 
| 237 | 
         
            +
                    null_embs: Optional[torch.FloatTensor] = None,
         
     | 
| 238 | 
         
            +
                    motion_field_strength: float = 12,
         
     | 
| 239 | 
         
            +
                    output_type: Optional[str] = "tensor",
         
     | 
| 240 | 
         
            +
                    return_dict: bool = True,
         
     | 
| 241 | 
         
            +
                    callback: Optional[Callable[[
         
     | 
| 242 | 
         
            +
                        int, int, torch.FloatTensor], None]] = None,
         
     | 
| 243 | 
         
            +
                    callback_steps: Optional[int] = 1,
         
     | 
| 244 | 
         
            +
                    use_foreground_motion_field: bool = True,
         
     | 
| 245 | 
         
            +
                    use_motion_field: bool = True,
         
     | 
| 246 | 
         
            +
                    smooth_bg: bool = True,
         
     | 
| 247 | 
         
            +
                    smooth_bg_strength: float = 0.4,
         
     | 
| 248 | 
         
            +
                    **kwargs,
         
     | 
| 249 | 
         
            +
                ):
         
     | 
| 250 | 
         
            +
             
     | 
| 251 | 
         
            +
                    print(f" Use: Motion field = {use_motion_field}")
         
     | 
| 252 | 
         
            +
                    print(f" Use: Background smoothing = {smooth_bg}")
         
     | 
| 253 | 
         
            +
                    # Default height and width to unet
         
     | 
| 254 | 
         
            +
                    height = height or self.unet.config.sample_size * self.vae_scale_factor
         
     | 
| 255 | 
         
            +
                    width = width or self.unet.config.sample_size * self.vae_scale_factor
         
     | 
| 256 | 
         
            +
             
     | 
| 257 | 
         
            +
                    # Check inputs. Raise error if not correct
         
     | 
| 258 | 
         
            +
                    self.check_inputs(prompt, height, width, callback_steps)
         
     | 
| 259 | 
         
            +
             
     | 
| 260 | 
         
            +
                    # Define call parameters
         
     | 
| 261 | 
         
            +
                    batch_size = 1 if isinstance(prompt, str) else len(prompt)
         
     | 
| 262 | 
         
            +
                    device = self._execution_device
         
     | 
| 263 | 
         
            +
                    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
         
     | 
| 264 | 
         
            +
                    # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
         
     | 
| 265 | 
         
            +
                    # corresponds to doing no classifier free guidance.
         
     | 
| 266 | 
         
            +
                    do_classifier_free_guidance = guidance_scale > 1.0
         
     | 
| 267 | 
         
            +
             
     | 
| 268 | 
         
            +
                    # Encode input prompt
         
     | 
| 269 | 
         
            +
                    text_embeddings = self._encode_prompt(
         
     | 
| 270 | 
         
            +
                        prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt
         
     | 
| 271 | 
         
            +
                    )
         
     | 
| 272 | 
         
            +
                    
         
     | 
| 273 | 
         
            +
                    # Prepare timesteps
         
     | 
| 274 | 
         
            +
                    self.scheduler.set_timesteps(num_inference_steps, device=device)
         
     | 
| 275 | 
         
            +
                    timesteps = self.scheduler.timesteps
         
     | 
| 276 | 
         
            +
                    
         
     | 
| 277 | 
         
            +
                    # print(f" Latent shape = {latents.shape}")
         
     | 
| 278 | 
         
            +
             
     | 
| 279 | 
         
            +
                    # Prepare latent variables
         
     | 
| 280 | 
         
            +
                    num_channels_latents = self.unet.in_channels
         
     | 
| 281 | 
         
            +
             
     | 
| 282 | 
         
            +
                    xT = self.prepare_latents(
         
     | 
| 283 | 
         
            +
                        batch_size * num_videos_per_prompt,
         
     | 
| 284 | 
         
            +
                        num_channels_latents,
         
     | 
| 285 | 
         
            +
                        video_length,
         
     | 
| 286 | 
         
            +
                        height,
         
     | 
| 287 | 
         
            +
                        width,
         
     | 
| 288 | 
         
            +
                        text_embeddings.dtype,
         
     | 
| 289 | 
         
            +
                        device,
         
     | 
| 290 | 
         
            +
                        generator,
         
     | 
| 291 | 
         
            +
                        xT,
         
     | 
| 292 | 
         
            +
                    )
         
     | 
| 293 | 
         
            +
                    dtype = xT.dtype
         
     | 
| 294 | 
         
            +
             
     | 
| 295 | 
         
            +
                    # when motion field is not used, augment with random latent codes
         
     | 
| 296 | 
         
            +
                    if use_motion_field:
         
     | 
| 297 | 
         
            +
                        xT = xT[:, :, :1]
         
     | 
| 298 | 
         
            +
                    else:
         
     | 
| 299 | 
         
            +
                        if xT.shape[2] < video_length:
         
     | 
| 300 | 
         
            +
                            xT_missing = self.prepare_latents(
         
     | 
| 301 | 
         
            +
                                batch_size * num_videos_per_prompt,
         
     | 
| 302 | 
         
            +
                                num_channels_latents,
         
     | 
| 303 | 
         
            +
                                video_length-xT.shape[2],
         
     | 
| 304 | 
         
            +
                                height,
         
     | 
| 305 | 
         
            +
                                width,
         
     | 
| 306 | 
         
            +
                                text_embeddings.dtype,
         
     | 
| 307 | 
         
            +
                                device,
         
     | 
| 308 | 
         
            +
                                generator,
         
     | 
| 309 | 
         
            +
                                None,
         
     | 
| 310 | 
         
            +
                            )
         
     | 
| 311 | 
         
            +
                            xT = torch.cat([xT, xT_missing], dim=2)
         
     | 
| 312 | 
         
            +
                    
         
     | 
| 313 | 
         
            +
             
     | 
| 314 | 
         
            +
                    xInit = xT.clone()
         
     | 
| 315 | 
         
            +
                    t0 = kwargs["t0"]
         
     | 
| 316 | 
         
            +
                    t1 = kwargs["t1"]
         
     | 
| 317 | 
         
            +
                    x_t1_1 = None
         
     | 
| 318 | 
         
            +
             
     | 
| 319 | 
         
            +
                    
         
     | 
| 320 | 
         
            +
                    # Prepare extra step kwargs.
         
     | 
| 321 | 
         
            +
                    extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
         
     | 
| 322 | 
         
            +
                    # Denoising loop
         
     | 
| 323 | 
         
            +
                    num_warmup_steps = len(timesteps) - \
         
     | 
| 324 | 
         
            +
                        num_inference_steps * self.scheduler.order
         
     | 
| 325 | 
         
            +
                   
         
     | 
| 326 | 
         
            +
             
     | 
| 327 | 
         
            +
             
     | 
| 328 | 
         
            +
                    ddim_res = self.DDIM_backward(num_inference_steps=num_inference_steps, timesteps=timesteps, skip_t=1000, t0=t0, t1=t1, do_classifier_free_guidance=do_classifier_free_guidance,
         
     | 
| 329 | 
         
            +
                                                      null_embs=null_embs, text_embeddings=text_embeddings, latents_local=xT, latents_dtype=dtype, guidance_scale=guidance_scale, guidance_stop_step=guidance_stop_step, callback=callback, callback_steps=callback_steps, extra_step_kwargs=extra_step_kwargs, num_warmup_steps=num_warmup_steps)
         
     | 
| 330 | 
         
            +
                    
         
     | 
| 331 | 
         
            +
                    x0 = ddim_res["x0"].detach()
         
     | 
| 332 | 
         
            +
                    
         
     | 
| 333 | 
         
            +
                    if "x_t0_1" in ddim_res:
         
     | 
| 334 | 
         
            +
                        x_t0_1 = ddim_res["x_t0_1"].detach()
         
     | 
| 335 | 
         
            +
                    if "x_t1_1" in ddim_res:
         
     | 
| 336 | 
         
            +
                        x_t1_1 = ddim_res["x_t1_1"].detach()
         
     | 
| 337 | 
         
            +
                    del ddim_res
         
     | 
| 338 | 
         
            +
                    del xT
         
     | 
| 339 | 
         
            +
             
     | 
| 340 | 
         
            +
                    if use_motion_field:
         
     | 
| 341 | 
         
            +
                        del x0
         
     | 
| 342 | 
         
            +
                        shape = (batch_size, num_channels_latents, 1, height //
         
     | 
| 343 | 
         
            +
                                 self.vae_scale_factor, width // self.vae_scale_factor)
         
     | 
| 344 | 
         
            +
                   
         
     | 
| 345 | 
         
            +
                   
         
     | 
| 346 | 
         
            +
                        x_t0_k = x_t0_1[:, :, :1, :, :].repeat(1, 1, video_length-1, 1, 1)
         
     | 
| 347 | 
         
            +
             
     | 
| 348 | 
         
            +
                        
         
     | 
| 349 | 
         
            +
                        reference_flow = torch.zeros(
         
     | 
| 350 | 
         
            +
                            (video_length-1, 2, 512, 512), device=x_t0_1.device, dtype=x_t0_1.dtype)
         
     | 
| 351 | 
         
            +
                        for fr_idx in range(video_length-1):
         
     | 
| 352 | 
         
            +
                            reference_flow[fr_idx, :, :, :] = motion_field_strength*(fr_idx+1)
         
     | 
| 353 | 
         
            +
             
     | 
| 354 | 
         
            +
                        for idx, latent in enumerate(x_t0_k):
         
     | 
| 355 | 
         
            +
                            x_t0_k[idx] = self.warp_latents_independently(
         
     | 
| 356 | 
         
            +
                                latent[None], reference_flow)
         
     | 
| 357 | 
         
            +
             
     | 
| 358 | 
         
            +
                        # assuming t0=t1=1000, if t0 = 1000
         
     | 
| 359 | 
         
            +
                        if t1 > t0:
         
     | 
| 360 | 
         
            +
                            x_t1_k = self.DDPM_forward(
         
     | 
| 361 | 
         
            +
                                x0=x_t0_k, t0=t0, tMax=t1, device=device, shape=shape, text_embeddings=text_embeddings, generator=generator)
         
     | 
| 362 | 
         
            +
                        else:
         
     | 
| 363 | 
         
            +
                            x_t1_k = x_t0_k
         
     | 
| 364 | 
         
            +
             
     | 
| 365 | 
         
            +
                        if x_t1_1 is None:
         
     | 
| 366 | 
         
            +
                            raise Exception
         
     | 
| 367 | 
         
            +
             
     | 
| 368 | 
         
            +
                        x_t1 = torch.cat([x_t1_1, x_t1_k], dim=2).clone().detach()
         
     | 
| 369 | 
         
            +
             
     | 
| 370 | 
         
            +
                        ddim_res = self.DDIM_backward(num_inference_steps=num_inference_steps, timesteps=timesteps, skip_t=t1, t0=-1, t1=-1, do_classifier_free_guidance=do_classifier_free_guidance,
         
     | 
| 371 | 
         
            +
                                                      null_embs=null_embs, text_embeddings=text_embeddings, latents_local=x_t1, latents_dtype=dtype, guidance_scale=guidance_scale, guidance_stop_step=guidance_stop_step, callback=callback, callback_steps=callback_steps, extra_step_kwargs=extra_step_kwargs, num_warmup_steps=num_warmup_steps)
         
     | 
| 372 | 
         
            +
             
     | 
| 373 | 
         
            +
                        x0 = ddim_res["x0"].detach()
         
     | 
| 374 | 
         
            +
                        del ddim_res
         
     | 
| 375 | 
         
            +
                    else:
         
     | 
| 376 | 
         
            +
                        x_t1 = x_t1_1.clone()
         
     | 
| 377 | 
         
            +
                        x_t1_1 = x_t1_1[:,:,:1,:,:].clone()
         
     | 
| 378 | 
         
            +
                        x_t1_k = x_t1_1[:,:,1:,:,:].clone()
         
     | 
| 379 | 
         
            +
                        x_t0_k = x_t0_1[:, :, 1:, :, :].clone()
         
     | 
| 380 | 
         
            +
                        x_t0_1 = x_t0_1[:,:,:1,:,:].clone()
         
     | 
| 381 | 
         
            +
             
     | 
| 382 | 
         
            +
             
     | 
| 383 | 
         
            +
                    move_object = use_foreground_motion_field
         
     | 
| 384 | 
         
            +
                    if move_object:
         
     | 
| 385 | 
         
            +
                        h, w = x0.shape[3], x0.shape[4]
         
     | 
| 386 | 
         
            +
                        # Move object
         
     | 
| 387 | 
         
            +
                        # reference_flow = torch.zeros(
         
     | 
| 388 | 
         
            +
                        #       (video_length-1, 2, 512, 512), device=x_t0_1.device, dtype=x_t0_1.dtype)
         
     | 
| 389 | 
         
            +
                        reference_flow_obj = torch.zeros(
         
     | 
| 390 | 
         
            +
                            (batch_size, video_length, 2, 512, 512), device=x_t0_1.device, dtype=x_t0_1.dtype)
         
     | 
| 391 | 
         
            +
             
     | 
| 392 | 
         
            +
                        for batch_idx, x0_b in enumerate(x0):
         
     | 
| 393 | 
         
            +
                            tmp = x0_b[None]
         
     | 
| 394 | 
         
            +
                            z0_b = []
         
     | 
| 395 | 
         
            +
                            for fr_split in range(tmp.shape[2]):
         
     | 
| 396 | 
         
            +
                                z0_b.append(self.decode_latents(
         
     | 
| 397 | 
         
            +
                                    tmp[:, :, fr_split, None]).detach())
         
     | 
| 398 | 
         
            +
                            z0_b = torch.cat(z0_b, dim=2)
         
     | 
| 399 | 
         
            +
                            z0_b = rearrange(z0_b[0], "c f h w -> f h w c")
         
     | 
| 400 | 
         
            +
                            shift = (-5 - 5) * torch.rand(2,
         
     | 
| 401 | 
         
            +
                                                          device=x0.device, dtype=x0.dtype) + 5
         
     | 
| 402 | 
         
            +
                            for frame_idx, z0_f in enumerate(z0_b):
         
     | 
| 403 | 
         
            +
                                if frame_idx > 0:
         
     | 
| 404 | 
         
            +
             
     | 
| 405 | 
         
            +
                                    z0_f = torch.round(
         
     | 
| 406 | 
         
            +
                                        z0_f * 255).cpu().numpy().astype(np.uint8)
         
     | 
| 407 | 
         
            +
                                    
         
     | 
| 408 | 
         
            +
                                    # apply SOD detection to obtain mask of foreground object
         
     | 
| 409 | 
         
            +
                                    m_f = torch.tensor(self.sod_model.process_data(
         
     | 
| 410 | 
         
            +
                                        z0_f), device=x0.device).to(x0.dtype)
         
     | 
| 411 | 
         
            +
                                    kernel = torch.ones(
         
     | 
| 412 | 
         
            +
                                        5, 5, device=x0.device, dtype=x0.dtype)
         
     | 
| 413 | 
         
            +
                                    mask = dilation(
         
     | 
| 414 | 
         
            +
                                        m_f[None, None].to(x0.device), kernel)[0]
         
     | 
| 415 | 
         
            +
                                    for coord_idx in range(2):
         
     | 
| 416 | 
         
            +
                                        reference_flow_obj[batch_idx, frame_idx,
         
     | 
| 417 | 
         
            +
                                                           coord_idx, :, :] = (1+frame_idx) * shift[coord_idx] * mask
         
     | 
| 418 | 
         
            +
             
     | 
| 419 | 
         
            +
             
     | 
| 420 | 
         
            +
             
     | 
| 421 | 
         
            +
                        for idx, x_t0_k_b in enumerate(x_t0_k):
         
     | 
| 422 | 
         
            +
                            x_t0_k[idx] = self.warp_latents_independently(
         
     | 
| 423 | 
         
            +
                                x_t0_k_b[None], reference_flow_obj[idx, 1:])
         
     | 
| 424 | 
         
            +
             
     | 
| 425 | 
         
            +
                        x_t1_k = self.DDPM_forward(
         
     | 
| 426 | 
         
            +
                            x0=x_t0_k, t0=t0, tMax=t1, device=device, shape=shape, text_embeddings=text_embeddings, generator=generator)
         
     | 
| 427 | 
         
            +
             
     | 
| 428 | 
         
            +
                        if x_t1_1 is None:
         
     | 
| 429 | 
         
            +
                            raise Exception
         
     | 
| 430 | 
         
            +
                        x_t1 = torch.cat([x_t1_1, x_t1_k], dim=2)
         
     | 
| 431 | 
         
            +
             
         
     | 
| 432 | 
         
            +
                        # del latent
         
     | 
| 433 | 
         
            +
                        ddim_res = self.DDIM_backward(num_inference_steps=num_inference_steps, timesteps=timesteps, skip_t=t1, t0=-1, t1=-1, do_classifier_free_guidance=do_classifier_free_guidance,
         
     | 
| 434 | 
         
            +
                                                      null_embs=null_embs, text_embeddings=text_embeddings, latents_local=x_t1, latents_dtype=dtype, guidance_scale=guidance_scale, guidance_stop_step=guidance_stop_step, callback=callback, callback_steps=callback_steps, extra_step_kwargs=extra_step_kwargs, num_warmup_steps=num_warmup_steps)
         
     | 
| 435 | 
         
            +
                        x0 = ddim_res["x0"].detach()
         
     | 
| 436 | 
         
            +
                        del ddim_res
         
     | 
| 437 | 
         
            +
             
     | 
| 438 | 
         
            +
             
     | 
| 439 | 
         
            +
                    # smooth background
         
     | 
| 440 | 
         
            +
                    if smooth_bg:
         
     | 
| 441 | 
         
            +
                        h, w = x0.shape[3], x0.shape[4]
         
     | 
| 442 | 
         
            +
                        M_FG = torch.zeros((batch_size, video_length, h, w),
         
     | 
| 443 | 
         
            +
                                           device=x0.device).to(x0.dtype)
         
     | 
| 444 | 
         
            +
                        for batch_idx, x0_b in enumerate(x0):
         
     | 
| 445 | 
         
            +
                            z0_b = self.decode_latents(x0_b[None]).detach()
         
     | 
| 446 | 
         
            +
                            z0_b = rearrange(z0_b[0], "c f h w -> f h w c")
         
     | 
| 447 | 
         
            +
                            for frame_idx, z0_f in enumerate(z0_b):
         
     | 
| 448 | 
         
            +
                                z0_f = torch.round(
         
     | 
| 449 | 
         
            +
                                    z0_f * 255).cpu().numpy().astype(np.uint8)
         
     | 
| 450 | 
         
            +
                                # apply SOD detection
         
     | 
| 451 | 
         
            +
                                m_f = torch.tensor(self.sod_model.process_data(
         
     | 
| 452 | 
         
            +
                                    z0_f), device=x0.device).to(x0.dtype)
         
     | 
| 453 | 
         
            +
                                mask = T.Resize(
         
     | 
| 454 | 
         
            +
                                    size=(h, w), interpolation=T.InterpolationMode.NEAREST)(m_f[None])
         
     | 
| 455 | 
         
            +
                                kernel = torch.ones(5, 5, device=x0.device, dtype=x0.dtype)
         
     | 
| 456 | 
         
            +
                                mask = dilation(mask[None].to(x0.device), kernel)[0]
         
     | 
| 457 | 
         
            +
                                M_FG[batch_idx, frame_idx, :, :] = mask
         
     | 
| 458 | 
         
            +
             
     | 
| 459 | 
         
            +
              
         
     | 
| 460 | 
         
            +
                        x_t1_1_fg_masked = x_t1_1 * \
         
     | 
| 461 | 
         
            +
                            (1 - repeat(M_FG[:, 0, :, :],
         
     | 
| 462 | 
         
            +
                                        "b w h -> b c 1 w h", c=x_t1_1.shape[1]))
         
     | 
| 463 | 
         
            +
             
     | 
| 464 | 
         
            +
             
     | 
| 465 | 
         
            +
                        x_t1_1_fg_masked_moved = []
         
     | 
| 466 | 
         
            +
                        for batch_idx, x_t1_1_fg_masked_b in enumerate(x_t1_1_fg_masked):
         
     | 
| 467 | 
         
            +
                            x_t1_fg_masked_b = x_t1_1_fg_masked_b.clone()
         
     | 
| 468 | 
         
            +
             
     | 
| 469 | 
         
            +
                            x_t1_fg_masked_b = x_t1_fg_masked_b.repeat(
         
     | 
| 470 | 
         
            +
                                1, video_length-1, 1, 1)
         
     | 
| 471 | 
         
            +
                            if use_motion_field:
         
     | 
| 472 | 
         
            +
                                x_t1_fg_masked_b = x_t1_fg_masked_b[None]
         
     | 
| 473 | 
         
            +
                                x_t1_fg_masked_b = self.warp_latents_independently(
         
     | 
| 474 | 
         
            +
                                    x_t1_fg_masked_b, reference_flow)
         
     | 
| 475 | 
         
            +
                            else:
         
     | 
| 476 | 
         
            +
                                x_t1_fg_masked_b = x_t1_fg_masked_b[None]
         
     | 
| 477 | 
         
            +
                            if move_object:
         
     | 
| 478 | 
         
            +
                                x_t1_fg_masked_b = self.warp_latents_independently(
         
     | 
| 479 | 
         
            +
                                    x_t1_fg_masked_b, reference_flow_obj[batch_idx, 1:])
         
     | 
| 480 | 
         
            +
             
     | 
| 481 | 
         
            +
                            x_t1_fg_masked_b = torch.cat(
         
     | 
| 482 | 
         
            +
                                [x_t1_1_fg_masked_b[None], x_t1_fg_masked_b], dim=2)
         
     | 
| 483 | 
         
            +
                            x_t1_1_fg_masked_moved.append(x_t1_fg_masked_b)
         
     | 
| 484 | 
         
            +
             
     | 
| 485 | 
         
            +
                        x_t1_1_fg_masked_moved = torch.cat(x_t1_1_fg_masked_moved, dim=0)
         
     | 
| 486 | 
         
            +
             
     | 
| 487 | 
         
            +
                        M_FG_1 = M_FG[:, :1, :, :]
         
     | 
| 488 | 
         
            +
             
     | 
| 489 | 
         
            +
                        M_FG_warped = []
         
     | 
| 490 | 
         
            +
                        for batch_idx, m_fg_1_b in enumerate(M_FG_1):
         
     | 
| 491 | 
         
            +
                            m_fg_1_b = m_fg_1_b[None, None]
         
     | 
| 492 | 
         
            +
                            m_fg_b = m_fg_1_b.repeat(1, 1, video_length-1, 1, 1)
         
     | 
| 493 | 
         
            +
                            if  use_motion_field:
         
     | 
| 494 | 
         
            +
                                m_fg_b = self.warp_latents_independently(
         
     | 
| 495 | 
         
            +
                                    m_fg_b.clone(), reference_flow)
         
     | 
| 496 | 
         
            +
                            if move_object:
         
     | 
| 497 | 
         
            +
                                m_fg_b = self.warp_latents_independently(
         
     | 
| 498 | 
         
            +
                                    m_fg_b, reference_flow_obj[batch_idx, 1:])
         
     | 
| 499 | 
         
            +
                            M_FG_warped.append(
         
     | 
| 500 | 
         
            +
                                torch.cat([m_fg_1_b[:1, 0], m_fg_b[:1, 0]], dim=1))
         
     | 
| 501 | 
         
            +
             
     | 
| 502 | 
         
            +
                        M_FG_warped = torch.cat(M_FG_warped, dim=0)
         
     | 
| 503 | 
         
            +
             
     | 
| 504 | 
         
            +
                        channels = x0.shape[1]
         
     | 
| 505 | 
         
            +
             
     | 
| 506 | 
         
            +
                        M_BG = (1-M_FG) * (1 - M_FG_warped)
         
     | 
| 507 | 
         
            +
                        M_BG = repeat(M_BG, "b f h w -> b c f h w", c=channels)
         
     | 
| 508 | 
         
            +
                        a_convex = smooth_bg_strength
         
     | 
| 509 | 
         
            +
              
         
     | 
| 510 | 
         
            +
                        x_t1_blending = (1-M_BG) * x_t1 + M_BG * (a_convex *
         
     | 
| 511 | 
         
            +
                                                                  x_t1 + (1-a_convex) * x_t1_1_fg_masked_moved)
         
     | 
| 512 | 
         
            +
             
     | 
| 513 | 
         
            +
                        '''
         
     | 
| 514 | 
         
            +
                        x_t1_blending = self.DDPM_forward(
         
     | 
| 515 | 
         
            +
                            x0=x_t1_blending, t0=t1, tMax=961, device=device, shape=shape, text_embeddings=text_embeddings, generator=generator)
         
     | 
| 516 | 
         
            +
                        t1 = 961
         
     | 
| 517 | 
         
            +
                        '''
         
     | 
| 518 | 
         
            +
                        latents = x_t1_blending
         
     | 
| 519 | 
         
            +
             
     | 
| 520 | 
         
            +
                        ddim_res = self.DDIM_backward(num_inference_steps=num_inference_steps, timesteps=timesteps, skip_t=t1, t0=-1, t1=-1, do_classifier_free_guidance=do_classifier_free_guidance,
         
     | 
| 521 | 
         
            +
                                                      null_embs=null_embs, text_embeddings=text_embeddings, latents_local=latents, latents_dtype=dtype, guidance_scale=guidance_scale, guidance_stop_step=guidance_stop_step, callback=callback, callback_steps=callback_steps, extra_step_kwargs=extra_step_kwargs, num_warmup_steps=num_warmup_steps)
         
     | 
| 522 | 
         
            +
                        x0 = ddim_res["x0"].detach()
         
     | 
| 523 | 
         
            +
                        del ddim_res
         
     | 
| 524 | 
         
            +
             
     | 
| 525 | 
         
            +
             
     | 
| 526 | 
         
            +
                    # Post-processing
         
     | 
| 527 | 
         
            +
                    video_list = []
         
     | 
| 528 | 
         
            +
                    for latent in x0:
         
     | 
| 529 | 
         
            +
                        tmp = latent[None]
         
     | 
| 530 | 
         
            +
                        print("Frame spit shape", tmp.shape)
         
     | 
| 531 | 
         
            +
                        frames = []
         
     | 
| 532 | 
         
            +
                        for fr_split in range(tmp.shape[2]):
         
     | 
| 533 | 
         
            +
                            print("frame decoding")
         
     | 
| 534 | 
         
            +
                            frames.append(self.decode_latents(
         
     | 
| 535 | 
         
            +
                                tmp[:, :, fr_split, None]).detach())
         
     | 
| 536 | 
         
            +
             
     | 
| 537 | 
         
            +
                        video_list.append(torch.cat(frames, dim=2).cpu().float().numpy())
         
     | 
| 538 | 
         
            +
             
     | 
| 539 | 
         
            +
                    # Convert to tensor
         
     | 
| 540 | 
         
            +
                    videos = []
         
     | 
| 541 | 
         
            +
                    if output_type == "tensor":
         
     | 
| 542 | 
         
            +
                        for video in video_list:
         
     | 
| 543 | 
         
            +
                            videos.append(torch.from_numpy(video))
         
     | 
| 544 | 
         
            +
                    if output_type == 'numpy':
         
     | 
| 545 | 
         
            +
                        for video in video_list:
         
     | 
| 546 | 
         
            +
                            videos.append(rearrange(video, 'b c f h w -> (b f) h w c'))
         
     | 
| 547 | 
         
            +
                    if not return_dict:
         
     | 
| 548 | 
         
            +
                        return video
         
     | 
| 549 | 
         
            +
             
     | 
| 550 | 
         
            +
                    return TextToVideoPipelineOutput(videos=videos, code=torch.split(xInit.detach().cpu(), 1, dim=0))
         
     |