Upload demo

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text

README.md

@@ -8,7 +8,7 @@ sdk_version: 5.23.1
 app_file: app.py
 pinned: false
 license: mit
-short_description: https://arxiv.org/abs/2501.08295
+short_description: "LayerAnimate: Layer-level Control for Animation"
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

__assets__/demos/demo_1/trajectory.json

@@ -0,0 +1,200 @@
+[
+  [
+    [
+      111.87965393066406,
+      204.28741455078125
+    ],
+    [
+      83.42483520507812,
+      204.21835327148438
+    ],
+    [
+      52.417137145996094,
+      205.34869384765625
+    ],
+    [
+      -10.01504135131836,
+      205.83694458007812
+    ],
+    [
+      -33.109561920166016,
+      206.53018188476562
+    ],
+    [
+      -86.02885437011719,
+      205.10772705078125
+    ],
+    [
+      -119.59435272216797,
+      204.4576873779297
+    ],
+    [
+      -168.70248413085938,
+      210.6188201904297
+    ],
+    [
+      -185.9542999267578,
+      211.16294860839844
+    ],
+    [
+      -206.82852172851562,
+      207.50912475585938
+    ],
+    [
+      -232.2637939453125,
+      208.35643005371094
+    ],
+    [
+      -177.6964111328125,
+      205.50949096679688
+    ],
+    [
+      -231.19761657714844,
+      203.8624267578125
+    ],
+    [
+      -276.06622314453125,
+      208.6024169921875
+    ],
+    [
+      -285.68218994140625,
+      210.30313110351562
+    ],
+    [
+      -235.0211639404297,
+      207.910400390625
+    ]
+  ],
+  [
+    [
+      130.59063720703125,
+      131.48106384277344
+    ],
+    [
+      101.31892395019531,
+      131.62567138671875
+    ],
+    [
+      69.3387451171875,
+      132.40696716308594
+    ],
+    [
+      6.821704864501953,
+      133.10546875
+    ],
+    [
+      -21.6120548248291,
+      132.92977905273438
+    ],
+    [
+      -83.36480712890625,
+      132.2947998046875
+    ],
+    [
+      -111.29481506347656,
+      131.91827392578125
+    ],
+    [
+      -168.74850463867188,
+      138.11587524414062
+    ],
+    [
+      -198.75299072265625,
+      139.32774353027344
+    ],
+    [
+      -253.08055114746094,
+      136.65480041503906
+    ],
+    [
+      -278.3507080078125,
+      136.42958068847656
+    ],
+    [
+      -312.9150390625,
+      134.22898864746094
+    ],
+    [
+      -332.20989990234375,
+      133.93161010742188
+    ],
+    [
+      -357.1211853027344,
+      139.33224487304688
+    ],
+    [
+      -361.4031677246094,
+      139.66172790527344
+    ],
+    [
+      -338.45501708984375,
+      141.38809204101562
+    ]
+  ],
+  [
+    [
+      308.344970703125,
+      6.6701483726501465
+    ],
+    [
+      278.66864013671875,
+      7.116205215454102
+    ],
+    [
+      247.65390014648438,
+      7.756659507751465
+    ],
+    [
+      184.76953125,
+      8.749884605407715
+    ],
+    [
+      154.9658203125,
+      8.66163444519043
+    ],
+    [
+      92.775146484375,
+      7.572597503662109
+    ],
+    [
+      63.20433044433594,
+      7.524573802947998
+    ],
+    [
+      1.4797935485839844,
+      13.07353401184082
+    ],
+    [
+      -26.288057327270508,
+      13.74260139465332
+    ],
+    [
+      -83.00379943847656,
+      11.522849082946777
+    ],
+    [
+      -109.52509307861328,
+      10.739717483520508
+    ],
+    [
+      -140.5462646484375,
+      8.596296310424805
+    ],
+    [
+      -155.35394287109375,
+      8.009984970092773
+    ],
+    [
+      -180.55775451660156,
+      13.584362030029297
+    ],
+    [
+      -185.0371856689453,
+      14.09956169128418
+    ],
+    [
+      -203.57778930664062,
+      18.082473754882812
+    ]
+  ]
+]

__assets__/demos/demo_1/trajectory.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:232a68740a9d2828277e786d760cb2d7436f4617ae1d64d31a61888be0c65ea1
+size 994

__assets__/demos/demo_2/trajectory.json

@@ -0,0 +1,200 @@
+[
+  [
+    [
+      158.21946716308594,
+      245.89105224609375
+    ],
+    [
+      148.94857788085938,
+      246.4789276123047
+    ],
+    [
+      137.88522338867188,
+      247.1299285888672
+    ],
+    [
+      128.4403839111328,
+      247.8033905029297
+    ],
+    [
+      127.84039306640625,
+      246.24864196777344
+    ],
+    [
+      127.06155395507812,
+      244.60606384277344
+    ],
+    [
+      126.77435302734375,
+      243.17208862304688
+    ],
+    [
+      126.42509460449219,
+      243.04747009277344
+    ],
+    [
+      125.61285400390625,
+      242.14913940429688
+    ],
+    [
+      125.40904235839844,
+      242.65948486328125
+    ],
+    [
+      125.03759765625,
+      242.90908813476562
+    ],
+    [
+      124.67877197265625,
+      242.95994567871094
+    ],
+    [
+      125.00759887695312,
+      242.61265563964844
+    ],
+    [
+      125.37916564941406,
+      242.13555908203125
+    ],
+    [
+      125.7420654296875,
+      242.410888671875
+    ],
+    [
+      125.54336547851562,
+      242.98825073242188
+    ]
+  ],
+  [
+    [
+      223.55435180664062,
+      204.28741455078125
+    ],
+    [
+      207.83377075195312,
+      202.7445068359375
+    ],
+    [
+      193.4696044921875,
+      200.418701171875
+    ],
+    [
+      178.7669677734375,
+      199.83621215820312
+    ],
+    [
+      178.14218139648438,
+      200.34848022460938
+    ],
+    [
+      176.58251953125,
+      200.19627380371094
+    ],
+    [
+      175.0523681640625,
+      200.24407958984375
+    ],
+    [
+      174.57379150390625,
+      199.90940856933594
+    ],
+    [
+      173.37542724609375,
+      200.4640350341797
+    ],
+    [
+      173.5262451171875,
+      200.5198974609375
+    ],
+    [
+      173.60935974121094,
+      200.36471557617188
+    ],
+    [
+      173.8643035888672,
+      200.39389038085938
+    ],
+    [
+      173.903076171875,
+      200.2958984375
+    ],
+    [
+      173.96859741210938,
+      200.00491333007812
+    ],
+    [
+      174.22422790527344,
+      200.09921264648438
+    ],
+    [
+      174.16683959960938,
+      200.00193786621094
+    ]
+  ],
+  [
+    [
+      232.88790893554688,
+      261.492431640625
+    ],
+    [
+      224.37376403808594,
+      258.9049072265625
+    ],
+    [
+      214.7504119873047,
+      255.82171630859375
+    ],
+    [
+      205.59695434570312,
+      252.74368286132812
+    ],
+    [
+      203.56024169921875,
+      254.83567810058594
+    ],
+    [
+      200.3128662109375,
+      256.933349609375
+    ],
+    [
+      197.56045532226562,
+      258.17236328125
+    ],
+    [
+      196.72007751464844,
+      258.3282470703125
+    ],
+    [
+      194.2041473388672,
+      259.42486572265625
+    ],
+    [
+      194.23858642578125,
+      259.9649353027344
+    ],
+    [
+      194.01547241210938,
+      260.14569091796875
+    ],
+    [
+      193.87156677246094,
+      259.9699401855469
+    ],
+    [
+      193.9617919921875,
+      259.7339172363281
+    ],
+    [
+      193.89659118652344,
+      259.5014343261719
+    ],
+    [
+      193.8680419921875,
+      259.7557373046875
+    ],
+    [
+      193.91842651367188,
+      260.28717041015625
+    ]
+  ]
+]

__assets__/demos/demo_2/trajectory.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba8194e3bd1376e10cb6c708d59603c406269b95bb1e266b20c7cfa66e248875
+size 972

__assets__/demos/demo_3/trajectory.json

@@ -0,0 +1,134 @@
+[
+  [
+    [
+      49.66927719116211,
+      126.28060150146484
+    ],
+    [
+      53.070796966552734,
+      140.00479125976562
+    ],
+    [
+      58.86982345581055,
+      157.8321533203125
+    ],
+    [
+      69.01676177978516,
+      175.84800720214844
+    ],
+    [
+      76.01651000976562,
+      197.62847900390625
+    ],
+    [
+      93.34223937988281,
+      232.17538452148438
+    ],
+    [
+      96.88280487060547,
+      246.68162536621094
+    ],
+    [
+      105.09373474121094,
+      265.91741943359375
+    ],
+    [
+      122.41947174072266,
+      300.46429443359375
+    ],
+    [
+      139.74520874023438,
+      335.0111999511719
+    ],
+    [
+      157.07093811035156,
+      369.55810546875
+    ],
+    [
+      174.39666748046875,
+      404.10498046875
+    ],
+    [
+      191.722412109375,
+      438.65185546875
+    ],
+    [
+      209.0481414794922,
+      473.19873046875
+    ],
+    [
+      226.37387084960938,
+      507.74560546875
+    ],
+    [
+      243.6995849609375,
+      542.29248046875
+    ]
+  ],
+  [
+    [
+      56.677669525146484,
+      69.07560729980469
+    ],
+    [
+      66.92218780517578,
+      90.37911224365234
+    ],
+    [
+      79.62323760986328,
+      116.14250183105469
+    ],
+    [
+      91.2628173828125,
+      141.8087921142578
+    ],
+    [
+      103.7956771850586,
+      167.58724975585938
+    ],
+    [
+      117.59683227539062,
+      195.22598266601562
+    ],
+    [
+      127.79037475585938,
+      221.12567138671875
+    ],
+    [
+      140.4638671875,
+      248.97164916992188
+    ],
+    [
+      138.9651641845703,
+      256.9488830566406
+    ],
+    [
+      165.24566650390625,
+      296.32525634765625
+    ],
+    [
+      191.52615356445312,
+      335.70166015625
+    ],
+    [
+      217.806640625,
+      375.07806396484375
+    ],
+    [
+      244.08714294433594,
+      414.4544372558594
+    ],
+    [
+      270.3676452636719,
+      453.830810546875
+    ],
+    [
+      296.64813232421875,
+      493.20721435546875
+    ],
+    [
+      322.92864990234375,
+      532.5836181640625
+    ]
+  ]
+]

__assets__/demos/demo_3/trajectory.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1080b8523b361f2e4fb3f5591c88f50e44d176a404e5f62b04cfc2bfe8c2f5d
+size 857

__assets__/demos/demo_4/trajectory.json

@@ -0,0 +1,200 @@
+[
+  [
+    [
+      186.72357177734375,
+      225.0892333984375
+    ],
+    [
+      186.59104919433594,
+      220.61599731445312
+    ],
+    [
+      190.39842224121094,
+      216.0291748046875
+    ],
+    [
+      199.52769470214844,
+      213.26031494140625
+    ],
+    [
+      204.145263671875,
+      214.56866455078125
+    ],
+    [
+      209.41751098632812,
+      214.23330688476562
+    ],
+    [
+      211.30255126953125,
+      216.12774658203125
+    ],
+    [
+      215.53131103515625,
+      215.55880737304688
+    ],
+    [
+      211.28453063964844,
+      215.3497314453125
+    ],
+    [
+      205.66819763183594,
+      210.34344482421875
+    ],
+    [
+      208.09231567382812,
+      197.720458984375
+    ],
+    [
+      201.51205444335938,
+      215.72598266601562
+    ],
+    [
+      191.19480895996094,
+      223.12850952148438
+    ],
+    [
+      194.90512084960938,
+      222.38108825683594
+    ],
+    [
+      200.74607849121094,
+      217.3187713623047
+    ],
+    [
+      207.563720703125,
+      235.63250732421875
+    ]
+  ],
+  [
+    [
+      289.63397216796875,
+      230.28970336914062
+    ],
+    [
+      289.8543701171875,
+      227.20205688476562
+    ],
+    [
+      292.2384033203125,
+      223.03854370117188
+    ],
+    [
+      301.47711181640625,
+      219.50289916992188
+    ],
+    [
+      308.8260803222656,
+      220.3004608154297
+    ],
+    [
+      315.6751403808594,
+      219.62095642089844
+    ],
+    [
+      317.8089599609375,
+      221.09295654296875
+    ],
+    [
+      320.73956298828125,
+      221.21011352539062
+    ],
+    [
+      317.1898193359375,
+      221.21250915527344
+    ],
+    [
+      319.5433349609375,
+      217.74606323242188
+    ],
+    [
+      317.6147155761719,
+      207.62603759765625
+    ],
+    [
+      308.29156494140625,
+      224.09878540039062
+    ],
+    [
+      294.7052917480469,
+      230.4814910888672
+    ],
+    [
+      298.7985534667969,
+      230.0016326904297
+    ],
+    [
+      304.0728454589844,
+      226.04998779296875
+    ],
+    [
+      314.6731872558594,
+      242.630126953125
+    ]
+  ],
+  [
+    [
+      214.7900390625,
+      230.28970336914062
+    ],
+    [
+      214.2034912109375,
+      226.12539672851562
+    ],
+    [
+      216.921630859375,
+      221.91062927246094
+    ],
+    [
+      226.7117156982422,
+      219.55148315429688
+    ],
+    [
+      232.1102294921875,
+      220.2542724609375
+    ],
+    [
+      237.49270629882812,
+      219.5577850341797
+    ],
+    [
+      240.1033935546875,
+      220.77169799804688
+    ],
+    [
+      243.27154541015625,
+      220.56069946289062
+    ],
+    [
+      240.3792724609375,
+      221.12344360351562
+    ],
+    [
+      235.10897827148438,
+      216.4136962890625
+    ],
+    [
+      234.0819091796875,
+      202.91900634765625
+    ],
+    [
+      224.08642578125,
+      220.4688720703125
+    ],
+    [
+      212.40911865234375,
+      227.7927703857422
+    ],
+    [
+      218.22300720214844,
+      226.47549438476562
+    ],
+    [
+      225.32315063476562,
+      221.8306884765625
+    ],
+    [
+      234.59808349609375,
+      239.94235229492188
+    ]
+  ]
+]

__assets__/demos/demo_4/trajectory.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c2904e38cbc8820daaa5f88085bbfc33aa3cd8b9be7d9588e02d6cadcccf2fa
+size 973

__assets__/demos/demo_5/trajectory.json

@@ -0,0 +1,332 @@
+[
+  [
+    [
+      494.2274169921875,
+      22.271512985229492
+    ],
+    [
+      499.44189453125,
+      21.746015548706055
+    ],
+    [
+      504.0919189453125,
+      21.225364685058594
+    ],
+    [
+      514.5880737304688,
+      20.82619285583496
+    ],
+    [
+      520.4939575195312,
+      20.672199249267578
+    ],
+    [
+      526.637451171875,
+      20.305557250976562
+    ],
+    [
+      534.9617919921875,
+      20.358591079711914
+    ],
+    [
+      539.2017211914062,
+      20.12591552734375
+    ],
+    [
+      543.9376220703125,
+      20.107173919677734
+    ],
+    [
+      549.5306396484375,
+      19.739456176757812
+    ],
+    [
+      553.4171142578125,
+      20.842308044433594
+    ],
+    [
+      554.49462890625,
+      20.15322494506836
+    ],
+    [
+      559.0555419921875,
+      21.292396545410156
+    ],
+    [
+      558.5130004882812,
+      21.357444763183594
+    ],
+    [
+      561.72607421875,
+      20.114139556884766
+    ],
+    [
+      560.4268798828125,
+      21.73964500427246
+    ]
+  ],
+  [
+    [
+      494.2274169921875,
+      48.27378463745117
+    ],
+    [
+      494.85711669921875,
+      48.05669403076172
+    ],
+    [
+      494.21563720703125,
+      48.0822868347168
+    ],
+    [
+      492.88446044921875,
+      48.20854187011719
+    ],
+    [
+      491.5914306640625,
+      48.36796569824219
+    ],
+    [
+      490.6370849609375,
+      48.649070739746094
+    ],
+    [
+      488.6202392578125,
+      48.874202728271484
+    ],
+    [
+      487.603271484375,
+      49.16374969482422
+    ],
+    [
+      486.469970703125,
+      49.414939880371094
+    ],
+    [
+      484.92120361328125,
+      49.98759460449219
+    ],
+    [
+      483.7000427246094,
+      50.26809310913086
+    ],
+    [
+      482.22125244140625,
+      50.42219161987305
+    ],
+    [
+      480.54931640625,
+      50.766448974609375
+    ],
+    [
+      479.24481201171875,
+      51.03229522705078
+    ],
+    [
+      478.1097106933594,
+      51.489837646484375
+    ],
+    [
+      476.470947265625,
+      52.048194885253906
+    ]
+  ],
+  [
+    [
+      64.8839111328125,
+      287.4947204589844
+    ],
+    [
+      81.71736145019531,
+      288.09869384765625
+    ],
+    [
+      100.02552795410156,
+      288.89111328125
+    ],
+    [
+      128.72686767578125,
+      289.8943176269531
+    ],
+    [
+      149.62322998046875,
+      290.7263488769531
+    ],
+    [
+      170.50192260742188,
+      291.29925537109375
+    ],
+    [
+      203.6192626953125,
+      292.2691345214844
+    ],
+    [
+      227.08547973632812,
+      292.68035888671875
+    ],
+    [
+      250.68621826171875,
+      293.3591613769531
+    ],
+    [
+      286.62176513671875,
+      294.1515197753906
+    ],
+    [
+      311.21240234375,
+      294.3829650878906
+    ],
+    [
+      335.68389892578125,
+      294.7114562988281
+    ],
+    [
+      373.18115234375,
+      295.2404479980469
+    ],
+    [
+      397.2961120605469,
+      295.111572265625
+    ],
+    [
+      422.346923828125,
+      295.5068054199219
+    ],
+    [
+      457.2431335449219,
+      295.49383544921875
+    ]
+  ],
+  [
+    [
+      64.8839111328125,
+      235.4901580810547
+    ],
+    [
+      61.33024597167969,
+      235.5504150390625
+    ],
+    [
+      57.36271667480469,
+      235.6099090576172
+    ],
+    [
+      50.592864990234375,
+      235.9037322998047
+    ],
+    [
+      46.184783935546875,
+      235.94981384277344
+    ],
+    [
+      42.2303466796875,
+      235.8488006591797
+    ],
+    [
+      35.333221435546875,
+      235.73272705078125
+    ],
+    [
+      29.864356994628906,
+      236.13253784179688
+    ],
+    [
+      24.596290588378906,
+      236.366943359375
+    ],
+    [
+      17.585124969482422,
+      236.61953735351562
+    ],
+    [
+      12.934989929199219,
+      236.7737274169922
+    ],
+    [
+      8.478790283203125,
+      236.75421142578125
+    ],
+    [
+      2.206012725830078,
+      236.9993896484375
+    ],
+    [
+      -2.862123489379883,
+      237.2617645263672
+    ],
+    [
+      -7.3507843017578125,
+      237.2784423828125
+    ],
+    [
+      -12.782325744628906,
+      237.2703094482422
+    ]
+  ],
+  [
+    [
+      92.88457489013672,
+      225.0892333984375
+    ],
+    [
+      88.737548828125,
+      225.09442138671875
+    ],
+    [
+      84.08223724365234,
+      225.36553955078125
+    ],
+    [
+      76.90846252441406,
+      225.7208251953125
+    ],
+    [
+      72.26066589355469,
+      225.9451141357422
+    ],
+    [
+      67.7042465209961,
+      226.13169860839844
+    ],
+    [
+      60.917144775390625,
+      226.32199096679688
+    ],
+    [
+      55.98236083984375,
+      226.5792236328125
+    ],
+    [
+      51.30162811279297,
+      226.9581298828125
+    ],
+    [
+      44.654823303222656,
+      227.06956481933594
+    ],
+    [
+      40.06951904296875,
+      227.15420532226562
+    ],
+    [
+      35.59206771850586,
+      227.13719177246094
+    ],
+    [
+      29.056011199951172,
+      227.17002868652344
+    ],
+    [
+      24.805736541748047,
+      227.24826049804688
+    ],
+    [
+      20.537612915039062,
+      227.34564208984375
+    ],
+    [
+      14.309333801269531,
+      227.30154418945312
+    ]
+  ]
+]

__assets__/demos/demo_5/trajectory.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e9da4a1142e8210f0486ff1682fe7853e8714ecf813bfef4b9019efbc102f61
+size 1222

app.py

@@ -0,0 +1,651 @@
+import argparse
+import datetime
+import os
+import json
+
+import torch
+import torchvision.transforms as transforms
+from torchvision.transforms import functional as F
+
+import spaces
+from huggingface_hub import snapshot_download
+import gradio as gr
+
+from diffusers import DDIMScheduler
+
+from lvdm.models.unet import UNetModel
+from lvdm.models.autoencoder import AutoencoderKL, AutoencoderKL_Dualref
+from lvdm.models.condition import FrozenOpenCLIPEmbedder, FrozenOpenCLIPImageEmbedderV2, Resampler
+from lvdm.models.layer_controlnet import LayerControlNet
+from lvdm.pipelines.pipeline_animation import AnimationPipeline
+from lvdm.utils import generate_gaussian_heatmap, save_videos_grid, save_videos_with_traj
+
+from einops import rearrange
+import cv2
+import decord
+from PIL import Image
+import numpy as np
+from scipy.interpolate import PchipInterpolator
+
+SAVE_DIR = "outputs"
+LENGTH = 16
+WIDTH = 512
+HEIGHT = 320
+LAYER_CAPACITY = 4
+DEVICE = "cuda"
+
+os.makedirs("checkpoints", exist_ok=True)
+
+snapshot_download(
+    "Yuppie1204/LayerAnimate-Mix",
+    local_dir="checkpoints/LayerAnimate-Mix",
+)
+
+class LayerAnimate:
+
+    @spaces.GPU
+    def __init__(self):
+        self.savedir = SAVE_DIR
+        os.makedirs(self.savedir, exist_ok=True)
+
+        self.weight_dtype  = torch.bfloat16
+        self.device        = DEVICE
+        self.text_encoder  = FrozenOpenCLIPEmbedder().eval()
+        self.image_encoder = FrozenOpenCLIPImageEmbedderV2().eval()
+
+        self.W = WIDTH
+        self.H = HEIGHT
+        self.L = LENGTH
+        self.layer_capacity = LAYER_CAPACITY
+
+        self.transforms = transforms.Compose([
+            transforms.Resize(min(self.H, self.W)),
+            transforms.CenterCrop((self.H, self.W)),
+        ])
+        self.pipeline = None
+        self.generator = None
+        # sample_grid is used to generate fixed trajectories to freeze static layers
+        self.sample_grid = np.meshgrid(np.linspace(0, self.W - 1, 10, dtype=int), np.linspace(0, self.H - 1, 10, dtype=int))
+        self.sample_grid = np.stack(self.sample_grid, axis=-1).reshape(-1, 1, 2)
+        self.sample_grid = np.repeat(self.sample_grid, self.L, axis=1) # [N, F, 2]
+
+    @spaces.GPU
+    def set_seed(self, seed):
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        self.generator = torch.Generator(self.device).manual_seed(seed)
+
+    @spaces.GPU
+    def set_model(self, pretrained_model_path):
+        scheduler = DDIMScheduler.from_pretrained(pretrained_model_path, subfolder="scheduler")
+        image_projector = Resampler.from_pretrained(pretrained_model_path, subfolder="image_projector").eval()
+        vae, vae_dualref = None, None
+        if "I2V" or "Mix" in pretrained_model_path:
+            vae           = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae").eval()
+        if "Interp" or "Mix" in pretrained_model_path:
+            vae_dualref   = AutoencoderKL_Dualref.from_pretrained(pretrained_model_path, subfolder="vae_dualref").eval()
+        unet              = UNetModel.from_pretrained(pretrained_model_path, subfolder="unet").eval()
+        layer_controlnet  = LayerControlNet.from_pretrained(pretrained_model_path, subfolder="layer_controlnet").eval()
+
+        self.pipeline = AnimationPipeline(
+            vae=vae, vae_dualref=vae_dualref, text_encoder=self.text_encoder, image_encoder=self.image_encoder, image_projector=image_projector,
+            unet=unet, layer_controlnet=layer_controlnet, scheduler=scheduler
+        ).to(device=self.device, dtype=self.weight_dtype)
+        if "Interp" or "Mix" in pretrained_model_path:
+            self.pipeline.vae_dualref.decoder.to(dtype=torch.float32)
+        return pretrained_model_path
+
+    def upload_image(self, image):
+        image = self.transforms(image)
+        return image
+
+    def run(self, input_image, input_image_end, pretrained_model_path, seed,
+            prompt, n_prompt, num_inference_steps, guidance_scale,
+            *layer_args):
+        self.set_seed(seed)
+        global layer_tracking_points
+        args_layer_tracking_points = [layer_tracking_points[i].value for i in range(self.layer_capacity)]
+
+        args_layer_masks = layer_args[:self.layer_capacity]
+        args_layer_masks_end = layer_args[self.layer_capacity : 2 * self.layer_capacity]
+        args_layer_controls = layer_args[2 * self.layer_capacity : 3 * self.layer_capacity]
+        args_layer_scores = list(layer_args[3 * self.layer_capacity : 4 * self.layer_capacity])
+        args_layer_sketches = layer_args[4 * self.layer_capacity : 5 * self.layer_capacity]
+        args_layer_valids = layer_args[5 * self.layer_capacity : 6 * self.layer_capacity]
+        args_layer_statics = layer_args[6 * self.layer_capacity : 7 * self.layer_capacity]
+        for layer_idx in range(self.layer_capacity):
+            if args_layer_controls[layer_idx] != "score":
+                args_layer_scores[layer_idx] = -1
+            if args_layer_statics[layer_idx]:
+                args_layer_scores[layer_idx] = 0
+
+        mode = "i2v"
+        image1 = F.to_tensor(input_image) * 2 - 1
+        frame_tensor = image1[None].to(self.device) # [F, C, H, W]
+        if input_image_end is not None:
+            mode = "interpolate"
+            image2 = F.to_tensor(input_image_end) * 2 - 1
+            frame_tensor2 = image2[None].to(self.device)
+            frame_tensor = torch.cat([frame_tensor, frame_tensor2], dim=0)
+        frame_tensor = frame_tensor[None]
+
+        if mode == "interpolate":
+            layer_masks = torch.zeros((1, self.layer_capacity, 2, 1, self.H, self.W), dtype=torch.bool)
+        else:
+            layer_masks = torch.zeros((1, self.layer_capacity, 1, 1, self.H, self.W), dtype=torch.bool)
+        for layer_idx in range(self.layer_capacity):
+            if args_layer_masks[layer_idx] is not None:
+                mask = F.to_tensor(args_layer_masks[layer_idx]) > 0.5
+                layer_masks[0, layer_idx, 0] = mask
+            if args_layer_masks_end[layer_idx] is not None and mode == "interpolate":
+                mask = F.to_tensor(args_layer_masks_end[layer_idx]) > 0.5
+                layer_masks[0, layer_idx, 1] = mask
+        layer_masks = layer_masks.to(self.device)
+        layer_regions = layer_masks * frame_tensor[:, None]
+        layer_validity = torch.tensor([args_layer_valids], dtype=torch.bool, device=self.device)
+        motion_scores = torch.tensor([args_layer_scores], dtype=self.weight_dtype, device=self.device)
+        layer_static = torch.tensor([args_layer_statics], dtype=torch.bool, device=self.device)
+
+        sketch = torch.ones((1, self.layer_capacity, self.L, 3, self.H, self.W), dtype=self.weight_dtype)
+        for layer_idx in range(self.layer_capacity):
+            sketch_path = args_layer_sketches[layer_idx]
+            if sketch_path is not None:
+                video_reader = decord.VideoReader(sketch_path)
+                assert len(video_reader) == self.L, f"Input the length of sketch sequence should match the video length."
+                video_frames = video_reader.get_batch(range(self.L)).asnumpy()
+                sketch_values = [F.to_tensor(self.transforms(Image.fromarray(frame))) for frame in video_frames]
+                sketch_values = torch.stack(sketch_values) * 2 - 1
+                sketch[0, layer_idx] = sketch_values
+        sketch = sketch.to(self.device)
+
+        heatmap = torch.zeros((1, self.layer_capacity, self.L, 3, self.H, self.W), dtype=self.weight_dtype)
+        heatmap[:, :, :, 0] -= 1
+        trajectory = []
+        traj_layer_index = []
+        for layer_idx in range(self.layer_capacity):
+            tracking_points = args_layer_tracking_points[layer_idx]
+            if args_layer_statics[layer_idx]:
+                # generate pseudo trajectory for static layers
+                temp_layer_mask = layer_masks[0, layer_idx, 0, 0].cpu().numpy()
+                valid_flag = temp_layer_mask[self.sample_grid[:, 0, 1], self.sample_grid[:, 0, 0]]
+                valid_grid = self.sample_grid[valid_flag]    # [F, N, 2]
+                trajectory.extend(list(valid_grid))
+                traj_layer_index.extend([layer_idx] * valid_grid.shape[0])
+            else:
+                for temp_track in tracking_points:
+                    if len(temp_track) > 1:
+                        x = [point[0] for point in temp_track]
+                        y = [point[1] for point in temp_track]
+                        t = np.linspace(0, 1, len(temp_track))
+                        fx = PchipInterpolator(t, x)
+                        fy = PchipInterpolator(t, y)
+                        t_new = np.linspace(0, 1, self.L)
+                        x_new = fx(t_new)
+                        y_new = fy(t_new)
+                        temp_traj = np.stack([x_new, y_new], axis=-1).astype(np.float32)
+                        trajectory.append(temp_traj)
+                        traj_layer_index.append(layer_idx)
+                    elif len(temp_track) == 1:
+                        trajectory.append(np.array(temp_track * self.L))
+                        traj_layer_index.append(layer_idx)
+
+        trajectory = np.stack(trajectory)
+        trajectory = np.transpose(trajectory, (1, 0, 2))
+        traj_layer_index = np.array(traj_layer_index)
+        heatmap = generate_gaussian_heatmap(trajectory, self.W, self.H, traj_layer_index, self.layer_capacity, offset=True)
+        heatmap = rearrange(heatmap, "f n c h w -> (f n) c h w")
+        graymap, offset = heatmap[:, :1], heatmap[:, 1:]
+        graymap = graymap / 255.
+        rad = torch.sqrt(offset[:, 0:1]**2 + offset[:, 1:2]**2)
+        rad_max = torch.max(rad)
+        epsilon = 1e-5
+        offset = offset / (rad_max + epsilon)
+        graymap = graymap * 2 - 1
+        heatmap = torch.cat([graymap, offset], dim=1)
+        heatmap = rearrange(heatmap, '(f n) c h w -> n f c h w', n=self.layer_capacity)
+        heatmap = heatmap[None]
+        heatmap = heatmap.to(self.device)
+
+        sample = self.pipeline(
+            prompt,
+            self.L,
+            self.H,
+            self.W,
+            frame_tensor,
+            layer_masks             = layer_masks,
+            layer_regions           = layer_regions,
+            layer_static            = layer_static,
+            motion_scores           = motion_scores,
+            sketch                  = sketch,
+            trajectory              = heatmap,
+            layer_validity          = layer_validity,
+            num_inference_steps     = num_inference_steps,
+            guidance_scale          = guidance_scale,
+            guidance_rescale        = 0.7,
+            negative_prompt         = n_prompt,
+            num_videos_per_prompt   = 1,
+            eta                     = 1.0,
+            generator               = self.generator,
+            fps                     = 24,
+            mode                    = mode,
+            weight_dtype            = self.weight_dtype,
+            output_type             = "tensor",
+        ).videos
+        output_video_path = os.path.join(self.savedir, "video.mp4")
+        save_videos_grid(sample, output_video_path, fps=8)
+        output_video_traj_path = os.path.join(self.savedir, "video_with_traj.mp4")
+        vis_traj_flag = np.zeros(trajectory.shape[1], dtype=bool)
+        for traj_idx in range(trajectory.shape[1]):
+            if not args_layer_statics[traj_layer_index[traj_idx]]:
+                vis_traj_flag[traj_idx] = True
+        vis_traj = torch.from_numpy(trajectory[:, vis_traj_flag])
+        save_videos_with_traj(sample[0], vis_traj, os.path.join(self.savedir, f"video_with_traj.mp4"), fps=8, line_width=7, circle_radius=10)
+        return output_video_path, output_video_traj_path
+
+
+def update_layer_region(image, layer_mask):
+    if image is None or layer_mask is None:
+        return None, False
+    layer_mask_tensor = (F.to_tensor(layer_mask) > 0.5).float()
+    image = F.to_tensor(image)
+    layer_region = image * layer_mask_tensor
+    layer_region = F.to_pil_image(layer_region)
+    layer_region.putalpha(layer_mask)
+    return layer_region, True
+
+def control_layers(control_type):
+    if control_type == "score":
+        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+    elif control_type == "trajectory":
+        return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)
+    else:
+        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+
+def visualize_trajectory(tracking_points, first_frame, first_mask, last_frame, last_mask):
+    first_mask_tensor = (F.to_tensor(first_mask) > 0.5).float()
+    first_frame = F.to_tensor(first_frame)
+    first_region = first_frame * first_mask_tensor
+    first_region = F.to_pil_image(first_region)
+    first_region.putalpha(first_mask)
+    transparent_background = first_region.convert('RGBA')
+
+    if last_frame is not None and last_mask is not None:
+        last_mask_tensor = (F.to_tensor(last_mask) > 0.5).float()
+        last_frame = F.to_tensor(last_frame)
+        last_region = last_frame * last_mask_tensor
+        last_region = F.to_pil_image(last_region)
+        last_region.putalpha(last_mask)
+        transparent_background_end = last_region.convert('RGBA')
+
+    width, height = transparent_background.size
+    transparent_layer = np.zeros((height, width, 4))
+    for track in tracking_points:
+        if len(track) > 1:
+            for i in range(len(track)-1):
+                start_point = np.array(track[i], dtype=np.int32)
+                end_point = np.array(track[i+1], dtype=np.int32)
+                vx = end_point[0] - start_point[0]
+                vy = end_point[1] - start_point[1]
+                arrow_length = max(np.sqrt(vx**2 + vy**2), 1)
+                if i == len(track)-2:
+                    cv2.arrowedLine(transparent_layer, tuple(start_point), tuple(end_point), (255, 0, 0, 255), 2, tipLength=8 / arrow_length)
+                else:
+                    cv2.line(transparent_layer, tuple(start_point), tuple(end_point), (255, 0, 0, 255), 2,)
+        elif len(track) == 1:
+            cv2.circle(transparent_layer, tuple(track[0]), 5, (255, 0, 0, 255), -1)
+    transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
+    trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
+    if last_frame is not None and last_mask is not None:
+        trajectory_map_end = Image.alpha_composite(transparent_background_end, transparent_layer)
+    else:
+        trajectory_map_end = None
+    return trajectory_map, trajectory_map_end
+
+def add_drag(layer_idx):
+    global layer_tracking_points
+    tracking_points = layer_tracking_points[layer_idx].value
+    tracking_points.append([])
+    return
+
+def delete_last_drag(layer_idx, first_frame, first_mask, last_frame, last_mask):
+    global layer_tracking_points
+    tracking_points = layer_tracking_points[layer_idx].value
+    tracking_points.pop()
+    trajectory_map, trajectory_map_end = visualize_trajectory(tracking_points, first_frame, first_mask, last_frame, last_mask)
+    return trajectory_map, trajectory_map_end
+
+def delete_last_step(layer_idx, first_frame, first_mask, last_frame, last_mask):
+    global layer_tracking_points
+    tracking_points = layer_tracking_points[layer_idx].value
+    tracking_points[-1].pop()
+    trajectory_map, trajectory_map_end = visualize_trajectory(tracking_points, first_frame, first_mask, last_frame, last_mask)
+    return trajectory_map, trajectory_map_end
+
+def add_tracking_points(layer_idx, first_frame, first_mask, last_frame, last_mask, evt: gr.SelectData):  # SelectData is a subclass of EventData
+    print(f"You selected {evt.value} at {evt.index} from {evt.target}")
+    global layer_tracking_points
+    tracking_points = layer_tracking_points[layer_idx].value
+    tracking_points[-1].append(evt.index)
+    trajectory_map, trajectory_map_end = visualize_trajectory(tracking_points, first_frame, first_mask, last_frame, last_mask)
+    return trajectory_map, trajectory_map_end
+
+def reset_states(layer_idx, first_frame, first_mask, last_frame, last_mask):
+    global layer_tracking_points
+    layer_tracking_points[layer_idx].value = [[]]
+    tracking_points = layer_tracking_points[layer_idx].value
+    trajectory_map, trajectory_map_end = visualize_trajectory(tracking_points, first_frame, first_mask, last_frame, last_mask)
+    return trajectory_map, trajectory_map_end
+
+def upload_tracking_points(tracking_path, layer_idx, first_frame, first_mask, last_frame, last_mask):
+    if tracking_path is None:
+        layer_region, _ = update_layer_region(first_frame, first_mask)
+        layer_region_end, _ = update_layer_region(last_frame, last_mask)
+        return layer_region, layer_region_end
+
+    global layer_tracking_points
+    with open(tracking_path, "r") as f:
+        tracking_points = json.load(f)
+    layer_tracking_points[layer_idx].value = tracking_points
+    trajectory_map, trajectory_map_end = visualize_trajectory(tracking_points, first_frame, first_mask, last_frame, last_mask)
+    return trajectory_map, trajectory_map_end
+
+def reset_all_controls():
+    global layer_tracking_points
+    outputs = []
+    # Reset tracking points states
+    for layer_idx in range(LAYER_CAPACITY):
+        layer_tracking_points[layer_idx].value = [[]]
+
+    # Reset global components
+    outputs.extend([
+        "an anime scene.",  # text prompt
+        "",                 # negative text prompt
+        50,                 # inference steps
+        7.5,                # guidance scale
+        42,                 # seed
+        None,               # input image
+        None,               # input image end
+        None,               # output video
+        None,               # output video with trajectory
+    ])
+    # Reset layer controls visibility
+    outputs.extend([None] * LAYER_CAPACITY)    # layer masks
+    outputs.extend([None] * LAYER_CAPACITY)    # layer masks end
+    outputs.extend([None] * LAYER_CAPACITY)    # layer regions
+    outputs.extend([None] * LAYER_CAPACITY)    # layer regions end
+    outputs.extend(["sketch"] * LAYER_CAPACITY)    # layer controls
+    outputs.extend([gr.update(visible=False, value=-1) for _ in range(LAYER_CAPACITY)])    # layer score controls
+    outputs.extend([gr.update(visible=False) for _ in range(4 * LAYER_CAPACITY)])    # layer trajectory control 4 buttons
+    outputs.extend([gr.update(visible=False, value=None) for _ in range(LAYER_CAPACITY)])    # layer trajectory file
+    outputs.extend([None] * LAYER_CAPACITY)    # layer sketch controls
+    outputs.extend([False] * LAYER_CAPACITY)    # layer validity
+    outputs.extend([False] * LAYER_CAPACITY)    # layer statics
+    return outputs
+
+if __name__ == "__main__":
+    with gr.Blocks() as demo:
+        gr.Markdown("""<h1 align="center">LayerAnimate: Layer-level Control for Animation</h1><br>""")
+
+        gr.Markdown("""Gradio Demo for <a href='https://arxiv.org/abs/2501.08295'><b>LayerAnimate: Layer-level Control for Animation</b></a>.<br>
+                    Github Repo can be found at https://github.com/IamCreateAI/LayerAnimate<br>
+                    The template is inspired by Framer.""")
+
+        gr.Image(label="LayerAnimate: Layer-level Control for Animation", value="__assets__/figs/demos.gif", height=540, width=960)
+
+        gr.Markdown("""## Usage: <br>
+                    1. Select a pretrained model via the "Pretrained Model" dropdown of choices in the right column.<br>
+                    2. Upload frames in the right column.<br>
+                    &ensp;  1.1.  Upload the first frame.<br>
+                    &ensp;  1.2. (Optional) Upload the last frame.<br>
+                    3. Input layer-level controls in the left column.<br>
+                    &ensp;  2.1. Upload layer mask images for each layer, which can be obtained from many tools such as https://huggingface.co/spaces/yumyum2081/SAM2-Image-Predictor.<br>
+                    &ensp;  2.2. Choose a control type from "motion score", "trajectory" and "sketch".<br>
+                    &ensp;  2.3. For trajectory control, you can draw trajectories on layer regions.<br>
+                    &ensp;  &ensp;  2.3.1. Click "Add New Trajectory" to add a new trajectory.<br>
+                    &ensp;  &ensp;  2.3.2. Click "Reset" to reset all trajectories.<br>
+                    &ensp;  &ensp;  2.3.3. Click "Delete Last Step" to delete the lastest clicked control point.<br>
+                    &ensp;  &ensp;  2.3.4. Click "Delete Last Trajectory" to delete the whole lastest path.<br>
+                    &ensp;  &ensp;  2.3.5. Or upload a trajectory file in json format, we provide examples below.<br>
+                    &ensp;  2.4. For sketch control, you can upload a sketch video.<br>
+                    4. We provide four layers for you to control, and it is not necessary to use all of them.<br>
+                    5. Click "Run" button to generate videos. <br>
+                    6. **Note: Remember to click "Clear" button to clear all the controls before switching to another example.**<br>
+                    """)
+
+        layeranimate = LayerAnimate()
+        layer_indices = [gr.Number(value=i, visible=False) for i in range(LAYER_CAPACITY)]
+        layer_tracking_points = [gr.State([[]]) for _ in range(LAYER_CAPACITY)]
+        layer_masks = []
+        layer_masks_end = []
+        layer_regions = []
+        layer_regions_end = []
+        layer_controls = []
+        layer_score_controls = []
+        layer_traj_controls = []
+        layer_traj_files = []
+        layer_sketch_controls = []
+        layer_statics = []
+        layer_valids = []
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                for layer_idx in range(LAYER_CAPACITY):
+                    with gr.Accordion(label=f"Layer {layer_idx+1}", open=True if layer_idx == 0 else False):
+                        gr.Markdown("""<div align="center"><b>Layer Masks</b></div>""")
+                        gr.Markdown("**Note**: Layer mask for the last frame is not required in I2V mode.")
+                        with gr.Row():
+                            with gr.Column():
+                                layer_masks.append(gr.Image(
+                                    label="Layer Mask for First Frame",
+                                    height=320,
+                                    width=512,
+                                    image_mode="L",
+                                    type="pil",
+                                ))
+
+                            with gr.Column():
+                                layer_masks_end.append(gr.Image(
+                                    label="Layer Mask for Last Frame",
+                                    height=320,
+                                    width=512,
+                                    image_mode="L",
+                                    type="pil",
+                                ))
+                        gr.Markdown("""<div align="center"><b>Layer Regions</b></div>""")
+                        with gr.Row():
+                            with gr.Column():
+                                layer_regions.append(gr.Image(
+                                    label="Layer Region for First Frame",
+                                    height=320,
+                                    width=512,
+                                    image_mode="RGBA",
+                                    type="pil",
+                                    # value=Image.new("RGBA", (512, 320), (255, 255, 255, 0)),
+                                ))
+
+                            with gr.Column():
+                                layer_regions_end.append(gr.Image(
+                                    label="Layer Region for Last Frame",
+                                    height=320,
+                                    width=512,
+                                    image_mode="RGBA",
+                                    type="pil",
+                                    # value=Image.new("RGBA", (512, 320), (255, 255, 255, 0)),
+                                ))
+                        layer_controls.append(
+                            gr.Radio(["score", "trajectory", "sketch"], label="Choose A Control Type", value="sketch")
+                        )
+                        layer_score_controls.append(
+                            gr.Number(label="Motion Score", value=-1, visible=False)
+                        )
+                        layer_traj_controls.append(
+                            [
+                                gr.Button(value="Add New Trajectory", visible=False),
+                                gr.Button(value="Reset", visible=False),
+                                gr.Button(value="Delete Last Step", visible=False),
+                                gr.Button(value="Delete Last Trajectory", visible=False),
+                            ]
+                        )
+                        layer_traj_files.append(
+                            gr.File(label="Trajectory File", visible=False)
+                        )
+                        layer_sketch_controls.append(
+                            gr.Video(label="Sketch", height=320, width=512, visible=True)
+                        )
+                        layer_controls[layer_idx].change(
+                            fn=control_layers,
+                            inputs=layer_controls[layer_idx],
+                            outputs=[layer_score_controls[layer_idx], *layer_traj_controls[layer_idx], layer_traj_files[layer_idx], layer_sketch_controls[layer_idx]]
+                        )
+                        with gr.Row():
+                            layer_valids.append(gr.Checkbox(label="Valid", info="Is the layer valid?"))
+                            layer_statics.append(gr.Checkbox(label="Static", info="Is the layer static?"))
+
+            with gr.Column(scale=1):
+                pretrained_model_path = gr.Dropdown(
+                    label="Pretrained Model",
+                    choices=[
+                        "None",
+                        "checkpoints/LayerAnimate-Mix",
+                    ],
+                    value="None",
+                )
+                text_prompt = gr.Textbox(label="Text Prompt", value="an anime scene.")
+                text_n_prompt = gr.Textbox(label="Negative Text Prompt", value="")
+                with gr.Row():
+                    num_inference_steps = gr.Number(label="Inference Steps", value=50, minimum=1, maximum=1000)
+                    guidance_scale = gr.Number(label="Guidance Scale", value=7.5)
+                    seed = gr.Number(label="Seed", value=42)
+                with gr.Row():
+                    input_image = gr.Image(
+                        label="First Frame",
+                        height=320,
+                        width=512,
+                        type="pil",
+                    )
+                    input_image_end = gr.Image(
+                        label="Last Frame",
+                        height=320,
+                        width=512,
+                        type="pil",
+                    )
+                run_button = gr.Button(value="Run")
+                with gr.Row():
+                    output_video = gr.Video(
+                        label="Output Video",
+                        height=320,
+                        width=512,
+                    )
+                    output_video_traj = gr.Video(
+                        label="Output Video with Trajectory",
+                        height=320,
+                        width=512,
+                    )
+                clear_button = gr.Button(value="Clear")
+
+        with gr.Row():
+            gr.Markdown("""
+                ## Citation
+                ```bibtex
+                @article{yang2025layeranimate,
+                author    = {Yang, Yuxue and Fan, Lue and Lin, Zuzeng and Wang, Feng and Zhang, Zhaoxiang},
+                title     = {LayerAnimate: Layer-level Control for Animation},
+                journal   = {arXiv preprint arXiv:2501.08295},
+                year      = {2025},
+                }
+                ```
+                """)
+
+        pretrained_model_path.input(layeranimate.set_model, pretrained_model_path, pretrained_model_path)
+        input_image.upload(layeranimate.upload_image, input_image, input_image)
+        input_image_end.upload(layeranimate.upload_image, input_image_end, input_image_end)
+        for i in range(LAYER_CAPACITY):
+            layer_masks[i].upload(layeranimate.upload_image, layer_masks[i], layer_masks[i])
+            layer_masks[i].change(update_layer_region, [input_image, layer_masks[i]], [layer_regions[i], layer_valids[i]])
+            layer_masks_end[i].upload(layeranimate.upload_image, layer_masks_end[i], layer_masks_end[i])
+            layer_masks_end[i].change(update_layer_region, [input_image_end, layer_masks_end[i]], [layer_regions_end[i], layer_valids[i]])
+            layer_traj_controls[i][0].click(add_drag, layer_indices[i], None)
+            layer_traj_controls[i][1].click(
+                reset_states,
+                [layer_indices[i], input_image, layer_masks[i], input_image_end, layer_masks_end[i]],
+                [layer_regions[i], layer_regions_end[i]]
+            )
+            layer_traj_controls[i][2].click(
+                delete_last_step,
+                [layer_indices[i], input_image, layer_masks[i], input_image_end, layer_masks_end[i]],
+                [layer_regions[i], layer_regions_end[i]]
+            )
+            layer_traj_controls[i][3].click(
+                delete_last_drag,
+                [layer_indices[i], input_image, layer_masks[i], input_image_end, layer_masks_end[i]],
+                [layer_regions[i], layer_regions_end[i]]
+            )
+            layer_traj_files[i].change(
+                upload_tracking_points,
+                [layer_traj_files[i], layer_indices[i], input_image, layer_masks[i], input_image_end, layer_masks_end[i]],
+                [layer_regions[i], layer_regions_end[i]]
+            )
+            layer_regions[i].select(
+                add_tracking_points,
+                [layer_indices[i], input_image, layer_masks[i], input_image_end, layer_masks_end[i]],
+                [layer_regions[i], layer_regions_end[i]]
+            )
+            layer_regions_end[i].select(
+                add_tracking_points,
+                [layer_indices[i], input_image, layer_masks[i], input_image_end, layer_masks_end[i]],
+                [layer_regions[i], layer_regions_end[i]]
+            )
+        run_button.click(
+            layeranimate.run,
+            [input_image, input_image_end, pretrained_model_path, seed, text_prompt, text_n_prompt, num_inference_steps, guidance_scale,
+             *layer_masks, *layer_masks_end, *layer_controls, *layer_score_controls, *layer_sketch_controls, *layer_valids, *layer_statics],
+            [output_video, output_video_traj]
+        )
+        clear_button.click(
+            reset_all_controls,
+            [],
+            [
+                text_prompt, text_n_prompt, num_inference_steps, guidance_scale, seed,
+                input_image, input_image_end, output_video, output_video_traj,
+                *layer_masks, *layer_masks_end, *layer_regions, *layer_regions_end,
+                *layer_controls, *layer_score_controls, *[button for temp_layer_controls in layer_traj_controls for button in temp_layer_controls], *layer_traj_files,
+                *layer_sketch_controls, *layer_valids, *layer_statics
+            ]
+        )
+        examples = gr.Examples(
+            examples=[
+                [
+                    "__assets__/demos/demo_3/first_frame.jpg", "__assets__/demos/demo_3/last_frame.jpg",
+                    "score",      "__assets__/demos/demo_3/layer_0.jpg", "__assets__/demos/demo_3/layer_0_last.jpg", 0.4, None,                                      None,                                 True, False,
+                    "score",      "__assets__/demos/demo_3/layer_1.jpg", "__assets__/demos/demo_3/layer_1_last.jpg", 0.2, None,                                      None,                                 True, False,
+                    "trajectory", "__assets__/demos/demo_3/layer_2.jpg", "__assets__/demos/demo_3/layer_2_last.jpg", -1,  "__assets__/demos/demo_3/trajectory.json", None,                                 True, False,
+                    "sketch",     "__assets__/demos/demo_3/layer_3.jpg", "__assets__/demos/demo_3/layer_3_last.jpg", -1,  None,                                      "__assets__/demos/demo_3/sketch.mp4", True, False,
+                    52
+                ],
+                [
+                    "__assets__/demos/demo_4/first_frame.jpg", None,
+                    "score",      "__assets__/demos/demo_4/layer_0.jpg", None, 0.0, None,                                      None,                                 True, True,
+                    "trajectory", "__assets__/demos/demo_4/layer_1.jpg", None, -1,  "__assets__/demos/demo_4/trajectory.json", None,                                 True, False,
+                    "sketch",     "__assets__/demos/demo_4/layer_2.jpg", None, -1,  None,                                      "__assets__/demos/demo_4/sketch.mp4", True, False,
+                    "score", None, None, -1, None, None, False, False,
+                    42
+                ],
+                [
+                    "__assets__/demos/demo_5/first_frame.jpg", None,
+                    "sketch",     "__assets__/demos/demo_5/layer_0.jpg", None, -1, None,                                      "__assets__/demos/demo_5/sketch.mp4", True, False,
+                    "trajectory", "__assets__/demos/demo_5/layer_1.jpg", None, -1, "__assets__/demos/demo_5/trajectory.json", None,                                 True, False,
+                    "score", None, None, -1, None, None, False, False,
+                    "score", None, None, -1, None, None, False, False,
+                    47
+                ],
+            ],
+            inputs=[
+                input_image, input_image_end,
+                layer_controls[0], layer_masks[0], layer_masks_end[0], layer_score_controls[0], layer_traj_files[0], layer_sketch_controls[0], layer_valids[0], layer_statics[0],
+                layer_controls[1], layer_masks[1], layer_masks_end[1], layer_score_controls[1], layer_traj_files[1], layer_sketch_controls[1], layer_valids[1], layer_statics[1],
+                layer_controls[2], layer_masks[2], layer_masks_end[2], layer_score_controls[2], layer_traj_files[2], layer_sketch_controls[2], layer_valids[2], layer_statics[2],
+                layer_controls[3], layer_masks[3], layer_masks_end[3], layer_score_controls[3], layer_traj_files[3], layer_sketch_controls[3], layer_valids[3], layer_statics[3],
+                seed
+            ],
+        )
+    demo.launch()

lvdm/basics.py

@@ -0,0 +1,100 @@
+# adopted from
+# https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
+# and
+# https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+# and
+# https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py
+#
+# thanks!
+
+import torch.nn as nn
+from .utils import instantiate_from_config
+
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+def scale_module(module, scale):
+    """
+    Scale the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().mul_(scale)
+    return module
+
+
+def conv_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D convolution module.
+    """
+    if dims == 1:
+        return nn.Conv1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.Conv2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.Conv3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def linear(*args, **kwargs):
+    """
+    Create a linear module.
+    """
+    return nn.Linear(*args, **kwargs)
+
+
+def avg_pool_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D average pooling module.
+    """
+    if dims == 1:
+        return nn.AvgPool1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.AvgPool2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.AvgPool3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def nonlinearity(type='silu'):
+    if type == 'silu':
+        return nn.SiLU()
+    elif type == 'leaky_relu':
+        return nn.LeakyReLU()
+
+
+class GroupNormSpecific(nn.GroupNorm):
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+
+
+def normalization(channels, num_groups=32):
+    """
+    Make a standard normalization layer.
+    :param channels: number of input channels.
+    :return: an nn.Module for normalization.
+    """
+    return GroupNormSpecific(num_groups, channels)
+
+
+class HybridConditioner(nn.Module):
+
+    def __init__(self, c_concat_config, c_crossattn_config):
+        super().__init__()
+        self.concat_conditioner = instantiate_from_config(c_concat_config)
+        self.crossattn_conditioner = instantiate_from_config(c_crossattn_config)
+
+    def forward(self, c_concat, c_crossattn):
+        c_concat = self.concat_conditioner(c_concat)
+        c_crossattn = self.crossattn_conditioner(c_crossattn)
+        return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}

lvdm/common.py

@@ -0,0 +1,94 @@
+import math
+from inspect import isfunction
+import torch
+from torch import nn
+import torch.distributed as dist
+
+
+def gather_data(data, return_np=True):
+    ''' gather data from multiple processes to one list '''
+    data_list = [torch.zeros_like(data) for _ in range(dist.get_world_size())]
+    dist.all_gather(data_list, data)  # gather not supported with NCCL
+    if return_np:
+        data_list = [data.cpu().numpy() for data in data_list]
+    return data_list
+
+def autocast(f):
+    def do_autocast(*args, **kwargs):
+        with torch.cuda.amp.autocast(enabled=True,
+                                     dtype=torch.get_autocast_gpu_dtype(),
+                                     cache_enabled=torch.is_autocast_cache_enabled()):
+            return f(*args, **kwargs)
+    return do_autocast
+
+
+def extract_into_tensor(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+
+def noise_like(shape, device, repeat=False):
+    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
+    noise = lambda: torch.randn(shape, device=device)
+    return repeat_noise() if repeat else noise()
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+def exists(val):
+    return val is not None
+
+def identity(*args, **kwargs):
+    return nn.Identity()
+
+def uniq(arr):
+    return{el: True for el in arr}.keys()
+
+def mean_flat(tensor):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+def ismap(x):
+    if not isinstance(x, torch.Tensor):
+        return False
+    return (len(x.shape) == 4) and (x.shape[1] > 3)
+
+def isimage(x):
+    if not isinstance(x,torch.Tensor):
+        return False
+    return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)
+
+def max_neg_value(t):
+    return -torch.finfo(t.dtype).max
+
+def shape_to_str(x):
+    shape_str = "x".join([str(x) for x in x.shape])
+    return shape_str
+
+def init_(tensor):
+    dim = tensor.shape[-1]
+    std = 1 / math.sqrt(dim)
+    tensor.uniform_(-std, std)
+    return tensor
+
+ckpt = torch.utils.checkpoint.checkpoint
+def checkpoint(func, inputs, params, flag):
+    """
+    Evaluate a function without caching intermediate activations, allowing for
+    reduced memory at the expense of extra compute in the backward pass.
+    :param func: the function to evaluate.
+    :param inputs: the argument sequence to pass to `func`.
+    :param params: a sequence of parameters `func` depends on but does not
+                   explicitly take as arguments.
+    :param flag: if False, disable gradient checkpointing.
+    """
+    if flag:
+        return ckpt(func, *inputs, use_reentrant=False)
+    else:
+        return func(*inputs)

lvdm/models/autoencoder.py

@@ -0,0 +1,143 @@
+import os
+from functools import partial
+from dataclasses import dataclass
+
+import torch
+import numpy as np
+from einops import rearrange
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models import ModelMixin
+from diffusers.utils import BaseOutput
+
+from ..modules.ae_modules import Encoder, Decoder
+from ..modules.ae_dualref_modules import VideoDecoder
+from ..utils import instantiate_from_config
+
+
+@dataclass
+class DecoderOutput(BaseOutput):
+    """
+    Output of decoding method.
+
+    Args:
+        sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            Decoded output sample of the model. Output of the last layer of the model.
+    """
+
+    sample: torch.FloatTensor
+
+
+@dataclass
+class AutoencoderKLOutput(BaseOutput):
+    """
+    Output of AutoencoderKL encoding method.
+
+    Args:
+        latent_dist (`DiagonalGaussianDistribution`):
+            Encoded outputs of `Encoder` represented as the mean and logvar of `DiagonalGaussianDistribution`.
+            `DiagonalGaussianDistribution` allows for sampling latents from the distribution.
+    """
+
+    latent_dist: "DiagonalGaussianDistribution"
+
+
+class AutoencoderKL(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(self,
+                 ddconfig,
+                 embed_dim,
+                 image_key="image",
+                 input_dim=4,
+                 use_checkpoint=False,
+                 ):
+        super().__init__()
+        self.image_key = image_key
+        self.encoder = Encoder(**ddconfig)
+        self.decoder = Decoder(**ddconfig)
+        assert ddconfig["double_z"]
+        self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)
+        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+        self.embed_dim = embed_dim
+        self.input_dim = input_dim
+        self.use_checkpoint = use_checkpoint
+
+    def encode(self, x, return_hidden_states=False, **kwargs):
+        if return_hidden_states:
+            h, hidden = self.encoder(x, return_hidden_states)
+            moments = self.quant_conv(h)
+            posterior = DiagonalGaussianDistribution(moments)
+            return AutoencoderKLOutput(latent_dist=posterior), hidden
+        else:
+            h = self.encoder(x)
+            moments = self.quant_conv(h)
+            posterior = DiagonalGaussianDistribution(moments)
+            return AutoencoderKLOutput(latent_dist=posterior)
+
+    def decode(self, z, **kwargs):
+        if len(kwargs) == 0: ## use the original decoder in AutoencoderKL
+            z = self.post_quant_conv(z)
+        dec = self.decoder(z, **kwargs)  ##change for SVD decoder by adding **kwargs
+        return dec
+
+    def forward(self, input, sample_posterior=True, **additional_decode_kwargs):
+        input_tuple = (input, )
+        forward_temp = partial(self._forward, sample_posterior=sample_posterior, **additional_decode_kwargs)
+        return checkpoint(forward_temp, input_tuple, self.parameters(), self.use_checkpoint)
+
+
+    def _forward(self, input, sample_posterior=True, **additional_decode_kwargs):
+        posterior = self.encode(input)[0]
+        if sample_posterior:
+            z = posterior.sample()
+        else:
+            z = posterior.mode()
+        dec = self.decode(z, **additional_decode_kwargs)
+        ## print(input.shape, dec.shape) torch.Size([16, 3, 256, 256]) torch.Size([16, 3, 256, 256])
+        return dec, posterior
+
+    def get_input(self, batch, k):
+        x = batch[k]
+        if x.dim() == 5 and self.input_dim == 4:
+            b,c,t,h,w = x.shape
+            self.b = b
+            self.t = t
+            x = rearrange(x, 'b c t h w -> (b t) c h w')
+
+        return x
+
+    def get_last_layer(self):
+        return self.decoder.conv_out.weight
+
+
+class AutoencoderKL_Dualref(AutoencoderKL):
+    @register_to_config
+    def __init__(self,
+                 ddconfig,
+                 embed_dim,
+                 image_key="image",
+                 input_dim=4,
+                 use_checkpoint=False,
+                 ):
+        super().__init__(ddconfig, embed_dim, image_key, input_dim, use_checkpoint)
+        self.decoder = VideoDecoder(**ddconfig)
+
+    def _forward(self, input, batch_size, sample_posterior=True, **additional_decode_kwargs):
+        posterior, hidden_states = self.encode(input, return_hidden_states=True)
+
+        hidden_states_first_last = []
+        ### use only the first and last hidden states
+        for hid in hidden_states:
+            hid = rearrange(hid, '(b t) c h w -> b c t h w', b=batch_size)
+            hid_new = torch.cat([hid[:, :, 0:1], hid[:, :, -1:]], dim=2)
+            hidden_states_first_last.append(hid_new)
+
+        if sample_posterior:
+            z = posterior[0].sample()
+        else:
+            z = posterior[0].mode()
+        dec = self.decode(z, ref_context=hidden_states_first_last, **additional_decode_kwargs)
+        ## print(input.shape, dec.shape) torch.Size([16, 3, 256, 256]) torch.Size([16, 3, 256, 256])
+        return dec, posterior

lvdm/models/condition.py

@@ -0,0 +1,477 @@
+import math
+import torch
+import torch.nn as nn
+from torchvision.transforms import functional as F
+import open_clip
+from torch.utils.checkpoint import checkpoint
+from transformers import T5Tokenizer, T5EncoderModel, CLIPTokenizer, CLIPTextModel
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models import ModelMixin
+from ..common import autocast
+from ..utils import count_params
+
+
+class AbstractEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def encode(self, *args, **kwargs):
+        raise NotImplementedError
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+class IdentityEncoder(AbstractEncoder):
+    def encode(self, x):
+        return x
+
+
+class ClassEmbedder(nn.Module):
+    def __init__(self, embed_dim, n_classes=1000, key='class', ucg_rate=0.1):
+        super().__init__()
+        self.key = key
+        self.embedding = nn.Embedding(n_classes, embed_dim)
+        self.n_classes = n_classes
+        self.ucg_rate = ucg_rate
+
+    def forward(self, batch, key=None, disable_dropout=False):
+        if key is None:
+            key = self.key
+        # this is for use in crossattn
+        c = batch[key][:, None]
+        if self.ucg_rate > 0. and not disable_dropout:
+            mask = 1. - torch.bernoulli(torch.ones_like(c) * self.ucg_rate)
+            c = mask * c + (1 - mask) * torch.ones_like(c) * (self.n_classes - 1)
+            c = c.long()
+        c = self.embedding(c)
+        return c
+
+    def get_unconditional_conditioning(self, bs, device="cuda"):
+        uc_class = self.n_classes - 1  # 1000 classes --> 0 ... 999, one extra class for ucg (class 1000)
+        uc = torch.ones((bs,), device=device) * uc_class
+        uc = {self.key: uc}
+        return uc
+
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+
+class FrozenT5Embedder(AbstractEncoder):
+    """Uses the T5 transformer encoder for text"""
+
+    def __init__(self, version="google/t5-v1_1-large", max_length=77,
+                 freeze=True):  # others are google/t5-v1_1-xl and google/t5-v1_1-xxl
+        super().__init__()
+        self.tokenizer = T5Tokenizer.from_pretrained(version)
+        self.transformer = T5EncoderModel.from_pretrained(version)
+        self.max_length = max_length  # TODO: typical value?
+        if freeze:
+            self.freeze()
+
+    def freeze(self):
+        self.transformer = self.transformer.eval()
+        # self.train = disabled_train
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, text):
+        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
+                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        tokens = batch_encoding["input_ids"].to(self.device)
+        outputs = self.transformer(input_ids=tokens)
+
+        z = outputs.last_hidden_state
+        return z
+
+    def encode(self, text):
+        return self(text)
+
+
+class FrozenCLIPEmbedder(AbstractEncoder):
+    """Uses the CLIP transformer encoder for text (from huggingface)"""
+    LAYERS = [
+        "last",
+        "pooled",
+        "hidden"
+    ]
+
+    def __init__(self, version="openai/clip-vit-large-patch14", max_length=77,
+                 freeze=True, layer="last", layer_idx=None):  # clip-vit-base-patch32
+        super().__init__()
+        assert layer in self.LAYERS
+        self.tokenizer = CLIPTokenizer.from_pretrained(version)
+        self.transformer = CLIPTextModel.from_pretrained(version)
+        self.max_length = max_length
+        if freeze:
+            self.freeze()
+        self.layer = layer
+        self.layer_idx = layer_idx
+        if layer == "hidden":
+            assert layer_idx is not None
+            assert 0 <= abs(layer_idx) <= 12
+
+    def freeze(self):
+        self.transformer = self.transformer.eval()
+        # self.train = disabled_train
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, text):
+        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
+                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        tokens = batch_encoding["input_ids"].to(self.device)
+        outputs = self.transformer(input_ids=tokens, output_hidden_states=self.layer == "hidden")
+        if self.layer == "last":
+            z = outputs.last_hidden_state
+        elif self.layer == "pooled":
+            z = outputs.pooler_output[:, None, :]
+        else:
+            z = outputs.hidden_states[self.layer_idx]
+        return z
+
+    def encode(self, text):
+        return self(text)
+
+
+class FrozenOpenCLIPEmbedder(AbstractEncoder):
+    """
+    Uses the OpenCLIP transformer encoder for text
+    """
+    LAYERS = [
+        # "pooled",
+        "last",
+        "penultimate"
+    ]
+
+    def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", max_length=77,
+                 freeze=True, layer="penultimate"):
+        super().__init__()
+        assert layer in self.LAYERS
+        model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), pretrained=version)
+        del model.visual
+        self.model = model
+
+        self.max_length = max_length
+        if freeze:
+            self.freeze()
+        self.layer = layer
+        if self.layer == "last":
+            self.layer_idx = 0
+        elif self.layer == "penultimate":
+            self.layer_idx = 1
+        else:
+            raise NotImplementedError()
+
+    def freeze(self):
+        self.model = self.model.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, text):
+        tokens = open_clip.tokenize(text) ## all clip models use 77 as context length
+        z = self.encode_with_transformer(tokens.to(self.device))
+        return z
+
+    def encode_with_transformer(self, text):
+        x = self.model.token_embedding(text)  # [batch_size, n_ctx, d_model]
+        x = x + self.model.positional_embedding
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.model.ln_final(x)
+        return x
+
+    def text_transformer_forward(self, x: torch.Tensor, attn_mask=None):
+        for i, r in enumerate(self.model.transformer.resblocks):
+            if i == len(self.model.transformer.resblocks) - self.layer_idx:
+                break
+            if self.model.transformer.grad_checkpointing and not torch.jit.is_scripting():
+                x = checkpoint(r, x, attn_mask)
+            else:
+                x = r(x, attn_mask=attn_mask)
+        return x
+
+    def encode(self, text):
+        return self(text)
+
+
+class FrozenOpenCLIPImageEmbedder(AbstractEncoder):
+    """
+    Uses the OpenCLIP vision transformer encoder for images
+    """
+
+    def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", max_length=77,
+                 freeze=True, layer="pooled", antialias=True, ucg_rate=0.):
+        super().__init__()
+        model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'),
+                                                            pretrained=version, )
+        del model.transformer
+        self.model = model
+        self.preprocess_val = preprocess_val
+        # self.mapper = torch.nn.Linear(1280, 1024)
+        self.max_length = max_length
+        if freeze:
+            self.freeze()
+        self.layer = layer
+        if self.layer == "penultimate":
+            raise NotImplementedError()
+            self.layer_idx = 1
+
+        self.antialias = antialias
+
+        self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
+        self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
+        self.ucg_rate = ucg_rate
+
+    def preprocess(self, x):
+        # normalize to [0,1]
+        x = F.resize(x, (224, 224), interpolation=F.InterpolationMode.BICUBIC, antialias=self.antialias)
+        x = (x + 1.) / 2.
+        # renormalize according to clip
+        x = F.normalize(x, mean=self.mean, std=self.std)
+        return x
+
+    def freeze(self):
+        self.model = self.model.eval()
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+    @autocast
+    def forward(self, image, no_dropout=False):
+        z = self.encode_with_vision_transformer(image)
+        if self.ucg_rate > 0. and not no_dropout:
+            z = torch.bernoulli((1. - self.ucg_rate) * torch.ones(z.shape[0], device=z.device))[:, None] * z
+        return z
+
+    def encode_with_vision_transformer(self, img):
+        img = self.preprocess(img)
+        x = self.model.visual(img)
+        return x
+
+    def encode(self, text):
+        return self(text)
+
+class FrozenOpenCLIPImageEmbedderV2(AbstractEncoder):
+    """
+    Uses the OpenCLIP vision transformer encoder for images
+    """
+
+    def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k",
+                 freeze=True, layer="pooled", antialias=True):
+        super().__init__()
+        model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'),
+                                                            pretrained=version, )
+        del model.transformer
+        self.model = model
+        self.preprocess_val = preprocess_val
+
+        if freeze:
+            self.freeze()
+        self.layer = layer
+        if self.layer == "penultimate":
+            raise NotImplementedError()
+            self.layer_idx = 1
+
+        self.antialias = antialias
+
+        self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
+        self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
+
+
+    def preprocess(self, x):
+        # normalize to [0,1]
+        x = F.resize(x, (224, 224), interpolation=F.InterpolationMode.BICUBIC, antialias=self.antialias)
+        x = (x + 1.) / 2.
+        # renormalize according to clip
+        x = F.normalize(x, mean=self.mean, std=self.std)
+        return x
+
+    def freeze(self):
+        self.model = self.model.eval()
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+    def forward(self, image, no_dropout=False):
+        ## image: b c h w
+        z = self.encode_with_vision_transformer(image)
+        return z
+
+    def encode_with_vision_transformer(self, x):
+        x = self.preprocess(x)
+
+        # to patches - whether to use dual patchnorm - https://arxiv.org/abs/2302.01327v1
+        if self.model.visual.input_patchnorm:
+            # einops - rearrange(x, 'b c (h p1) (w p2) -> b (h w) (c p1 p2)')
+            x = x.reshape(x.shape[0], x.shape[1], self.model.visual.grid_size[0], self.model.visual.patch_size[0], self.model.visual.grid_size[1], self.model.visual.patch_size[1])
+            x = x.permute(0, 2, 4, 1, 3, 5)
+            x = x.reshape(x.shape[0], self.model.visual.grid_size[0] * self.model.visual.grid_size[1], -1)
+            x = self.model.visual.patchnorm_pre_ln(x)
+            x = self.model.visual.conv1(x)
+        else:
+            x = self.model.visual.conv1(x)  # shape = [*, width, grid, grid]
+            x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+            x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+
+        # class embeddings and positional embeddings
+        x = torch.cat(
+            [self.model.visual.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device),
+             x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.model.visual.positional_embedding.to(x.dtype)
+
+        # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
+        x = self.model.visual.patch_dropout(x)
+        x = self.model.visual.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.model.visual.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        return x
+
+class FrozenCLIPT5Encoder(AbstractEncoder):
+    def __init__(self, clip_version="openai/clip-vit-large-patch14", t5_version="google/t5-v1_1-xl",
+                 clip_max_length=77, t5_max_length=77):
+        super().__init__()
+        self.clip_encoder = FrozenCLIPEmbedder(clip_version, max_length=clip_max_length)
+        self.t5_encoder = FrozenT5Embedder(t5_version, max_length=t5_max_length)
+        print(f"{self.clip_encoder.__class__.__name__} has {count_params(self.clip_encoder) * 1.e-6:.2f} M parameters, "
+              f"{self.t5_encoder.__class__.__name__} comes with {count_params(self.t5_encoder) * 1.e-6:.2f} M params.")
+
+    def encode(self, text):
+        return self(text)
+
+    def forward(self, text):
+        clip_z = self.clip_encoder.encode(text)
+        t5_z = self.t5_encoder.encode(text)
+        return [clip_z, t5_z]
+
+
+# FFN
+def FeedForward(dim, mult=4):
+    inner_dim = int(dim * mult)
+    return nn.Sequential(
+        nn.LayerNorm(dim),
+        nn.Linear(dim, inner_dim, bias=False),
+        nn.GELU(),
+        nn.Linear(inner_dim, dim, bias=False),
+    )
+
+
+def reshape_tensor(x, heads):
+    bs, length, width = x.shape
+    #(bs, length, width) --> (bs, length, n_heads, dim_per_head)
+    x = x.view(bs, length, heads, -1)
+    # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
+    x = x.transpose(1, 2)
+    # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
+    x = x.reshape(bs, heads, length, -1)
+    return x
+
+
+class PerceiverAttention(nn.Module):
+    def __init__(self, *, dim, dim_head=64, heads=8):
+        super().__init__()
+        self.scale = dim_head**-0.5
+        self.dim_head = dim_head
+        self.heads = heads
+        inner_dim = dim_head * heads
+
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
+        self.to_out = nn.Linear(inner_dim, dim, bias=False)
+
+
+    def forward(self, x, latents):
+        """
+        Args:
+            x (torch.Tensor): image features
+                shape (b, n1, D)
+            latent (torch.Tensor): latent features
+                shape (b, n2, D)
+        """
+        x = self.norm1(x)
+        latents = self.norm2(latents)
+
+        b, l, _ = latents.shape
+
+        q = self.to_q(latents)
+        kv_input = torch.cat((x, latents), dim=-2)
+        k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+
+        q = reshape_tensor(q, self.heads)
+        k = reshape_tensor(k, self.heads)
+        v = reshape_tensor(v, self.heads)
+
+        # attention
+        scale = 1 / math.sqrt(math.sqrt(self.dim_head))
+        weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        out = weight @ v
+
+        out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
+
+        return self.to_out(out)
+
+
+class Resampler(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+        self,
+        dim=1024,
+        depth=8,
+        dim_head=64,
+        heads=16,
+        num_queries=8,
+        embedding_dim=768,
+        output_dim=1024,
+        ff_mult=4,
+        video_length=None, # using frame-wise version or not
+    ):
+        super().__init__()
+        ## queries for a single frame / image
+        self.num_queries = num_queries
+        self.video_length = video_length
+
+        ## <num_queries> queries for each frame
+        if video_length is not None:
+            num_queries = num_queries * video_length
+
+        self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5)
+        self.proj_in = nn.Linear(embedding_dim, dim)
+        self.proj_out = nn.Linear(dim, output_dim)
+        self.norm_out = nn.LayerNorm(output_dim)
+
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
+                        FeedForward(dim=dim, mult=ff_mult),
+                    ]
+                )
+            )
+
+    def forward(self, x):
+        latents = self.latents.repeat(x.size(0), 1, 1) ## B (T L) C
+        x = self.proj_in(x)
+
+        for attn, ff in self.layers:
+            latents = attn(x, latents) + latents
+            latents = ff(latents) + latents
+
+        latents = self.proj_out(latents)
+        latents = self.norm_out(latents) # B L C or B (T L) C
+
+        return latents

lvdm/models/controlnet.py

@@ -0,0 +1,500 @@
+from typing import Any, Dict, List, Optional, Tuple, Union
+from einops import rearrange, repeat
+import numpy as np
+from functools import partial
+import torch
+from torch import nn
+from torch.nn import functional as F
+from .unet import TimestepEmbedSequential, ResBlock, Downsample, Upsample, TemporalConvBlock
+from ..basics import zero_module, conv_nd
+from ..modules.attention import SpatialTransformer, TemporalTransformer
+from ..common import checkpoint
+
+from diffusers import __version__
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.model_loading_utils import load_state_dict
+from diffusers.utils import (
+    SAFETENSORS_WEIGHTS_NAME,
+    WEIGHTS_NAME,
+    logging,
+    _get_model_file,
+    _add_variant
+)
+from omegaconf import ListConfig, DictConfig, OmegaConf
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class ResBlock_v2(nn.Module):
+    def __init__(
+        self,
+        channels,
+        emb_channels,
+        dropout,
+        out_channels=None,
+        dims=2,
+        use_checkpoint=False,
+        use_conv=False,
+        up=False,
+        down=False,
+        use_temporal_conv=False,
+        tempspatial_aware=False
+    ):
+        super().__init__()
+        self.channels = channels
+        self.emb_channels = emb_channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_checkpoint = use_checkpoint
+        self.use_temporal_conv = use_temporal_conv
+
+        self.in_layers = nn.Sequential(
+            nn.GroupNorm(32, channels),
+            nn.SiLU(),
+            zero_module(conv_nd(dims, channels, self.out_channels, 3, padding=1)),
+        )
+
+        self.updown = up or down
+
+        if up:
+            self.h_upd = Upsample(channels, False, dims)
+            self.x_upd = Upsample(channels, False, dims)
+        elif down:
+            self.h_upd = Downsample(channels, False, dims)
+            self.x_upd = Downsample(channels, False, dims)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 3, padding=1)
+        else:
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
+
+        if self.use_temporal_conv:
+            self.temopral_conv = TemporalConvBlock(
+                self.out_channels,
+                self.out_channels,
+                dropout=0.1,
+                spatial_aware=tempspatial_aware
+            )
+
+    def forward(self, x, batch_size=None):
+        """
+        Apply the block to a Tensor, conditioned on a timestep embedding.
+        :param x: an [N x C x ...] Tensor of features.
+        :return: an [N x C x ...] Tensor of outputs.
+        """
+        input_tuple = (x, )
+        if batch_size:
+            forward_batchsize = partial(self._forward, batch_size=batch_size)
+            return checkpoint(forward_batchsize, input_tuple, self.parameters(), self.use_checkpoint)
+        return checkpoint(self._forward, input_tuple, self.parameters(), self.use_checkpoint)
+
+    def _forward(self, x, batch_size=None):
+        if self.updown:
+            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+            h = in_rest(x)
+            h = self.h_upd(h)
+            x = self.x_upd(x)
+            h = in_conv(h)
+        else:
+            h = self.in_layers(x)
+        h = self.skip_connection(x) + h
+
+        if self.use_temporal_conv and batch_size:
+            h = rearrange(h, '(b t) c h w -> b c t h w', b=batch_size)
+            h = self.temopral_conv(h)
+            h = rearrange(h, 'b c t h w -> (b t) c h w')
+        return h
+
+
+class TrajectoryEncoder(nn.Module):
+    def __init__(self, cin, time_embed_dim, channels=[320, 640, 1280, 1280], nums_rb=3,
+                 dropout=0.0, use_checkpoint=False, tempspatial_aware=False, temporal_conv=False):
+        super(TrajectoryEncoder, self).__init__()
+        # self.unshuffle = nn.PixelUnshuffle(8)
+        self.channels = channels
+        self.nums_rb = nums_rb
+        self.body = []
+        # self.conv_out = []
+        for i in range(len(channels)):
+            for j in range(nums_rb):
+                if (i != 0) and (j == 0):
+                    self.body.append(
+                        ResBlock_v2(channels[i - 1], time_embed_dim, dropout,
+                            out_channels=channels[i], dims=2, use_checkpoint=use_checkpoint,
+                            tempspatial_aware=tempspatial_aware,
+                            use_temporal_conv=temporal_conv,
+                            down=True
+                        )
+                    )
+                else:
+                    self.body.append(
+                        ResBlock_v2(channels[i], time_embed_dim, dropout,
+                            out_channels=channels[i], dims=2, use_checkpoint=use_checkpoint,
+                            tempspatial_aware=tempspatial_aware,
+                            use_temporal_conv=temporal_conv,
+                            down=False
+                        )
+                    )
+        self.body.append(
+            ResBlock_v2(channels[-1], time_embed_dim, dropout,
+                out_channels=channels[-1], dims=2, use_checkpoint=use_checkpoint,
+                tempspatial_aware=tempspatial_aware,
+                use_temporal_conv=temporal_conv,
+                down=True
+            )
+        )
+        self.body = nn.ModuleList(self.body)
+        self.conv_in = nn.Conv2d(cin, channels[0], 3, 1, 1)
+        self.conv_out = zero_module(conv_nd(2, channels[-1], channels[-1], 3, 1, 1))
+
+    def forward(self, x, batch_size=None):
+        # unshuffle
+        # x = self.unshuffle(x)
+        # extract features
+        # features = []
+        x = self.conv_in(x)
+        for i in range(len(self.channels)):
+            for j in range(self.nums_rb):
+                idx = i * self.nums_rb + j
+                x = self.body[idx](x, batch_size)
+        x = self.body[-1](x, batch_size)
+        out = self.conv_out(x)
+        return out
+
+
+class ControlNet(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels,
+        model_channels,
+        out_channels,
+        num_res_blocks,
+        attention_resolutions,
+        dropout=0.0,
+        channel_mult=(1, 2, 4, 8),
+        conv_resample=True,
+        dims=2,
+        context_dim=None,
+        use_scale_shift_norm=False,
+        resblock_updown=False,
+        num_heads=-1,
+        num_head_channels=-1,
+        transformer_depth=1,
+        use_linear=False,
+        use_checkpoint=False,
+        temporal_conv=False,
+        tempspatial_aware=False,
+        temporal_attention=True,
+        use_relative_position=True,
+        use_causal_attention=False,
+        temporal_length=None,
+        addition_attention=False,
+        temporal_selfatt_only=True,
+        image_cross_attention=False,
+        image_cross_attention_scale_learnable=False,
+        default_fps=4,
+        fps_condition=False,
+        ignore_noisy_latents=True,
+        conditioning_channels=4,
+    ):
+        super().__init__()
+        if num_heads == -1:
+            assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
+        if num_head_channels == -1:
+            assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
+
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.temporal_attention = temporal_attention
+        time_embed_dim = model_channels * 4
+        self.use_checkpoint = use_checkpoint
+        temporal_self_att_only = True
+        self.addition_attention = addition_attention
+        self.temporal_length = temporal_length
+        self.image_cross_attention = image_cross_attention
+        self.image_cross_attention_scale_learnable = image_cross_attention_scale_learnable
+        self.default_fps = default_fps
+        self.fps_condition = fps_condition
+        self.ignore_noisy_latents = ignore_noisy_latents
+
+        ## Time embedding blocks
+        self.time_proj = Timesteps(model_channels, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.time_embed = TimestepEmbedding(model_channels, time_embed_dim)
+
+        if fps_condition:
+            self.fps_embedding = TimestepEmbedding(model_channels, time_embed_dim)
+            nn.init.zeros_(self.fps_embedding.linear_2.weight)
+            nn.init.zeros_(self.fps_embedding.linear_2.bias)
+
+        # self.cond_embedding = TrajectoryEncoder(
+        #         cin=conditioning_channels, time_embed_dim=time_embed_dim, channels=trajectory_channels, nums_rb=3,
+        #         dropout=dropout, use_checkpoint=use_checkpoint, tempspatial_aware=tempspatial_aware, temporal_conv=False
+        #     )
+        self.cond_embedding = zero_module(conv_nd(dims, conditioning_channels, model_channels, 3, padding=1))
+        self.input_blocks = nn.ModuleList(
+            [
+                TimestepEmbedSequential(conv_nd(dims, in_channels, model_channels, 3, padding=1))
+            ]
+        )
+
+        ## Output Block
+        self.downsample_output = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.GroupNorm(32, model_channels),
+                    nn.SiLU(),
+                    zero_module(conv_nd(dims, model_channels, model_channels, 3, padding=1))
+                )
+            ]
+        )
+
+        if self.addition_attention:
+            self.init_attn = TimestepEmbedSequential(
+                TemporalTransformer(
+                    model_channels,
+                    n_heads=8,
+                    d_head=num_head_channels,
+                    depth=transformer_depth,
+                    context_dim=context_dim,
+                    use_checkpoint=use_checkpoint, only_self_att=temporal_selfatt_only,
+                    causal_attention=False, relative_position=use_relative_position,
+                    temporal_length=temporal_length
+                )
+            )
+
+        ch = model_channels
+        ds = 1
+        for level, mult in enumerate(channel_mult):
+            for _ in range(num_res_blocks):
+                layers = [
+                    ResBlock(ch, time_embed_dim, dropout,
+                        out_channels=mult * model_channels, dims=dims, use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm, tempspatial_aware=tempspatial_aware,
+                        use_temporal_conv=temporal_conv
+                    )
+                ]
+                ch = mult * model_channels
+                if ds in attention_resolutions:
+                    if num_head_channels == -1:
+                        dim_head = ch // num_heads
+                    else:
+                        num_heads = ch // num_head_channels
+                        dim_head = num_head_channels
+                    layers.append(
+                        SpatialTransformer(ch, num_heads, dim_head,
+                            depth=transformer_depth, context_dim=context_dim, use_linear=use_linear,
+                            use_checkpoint=use_checkpoint, disable_self_attn=False,
+                            video_length=temporal_length, image_cross_attention=self.image_cross_attention,
+                            image_cross_attention_scale_learnable=self.image_cross_attention_scale_learnable,
+                        )
+                    )
+                    if self.temporal_attention:
+                        layers.append(
+                            TemporalTransformer(ch, num_heads, dim_head,
+                                depth=transformer_depth, context_dim=context_dim, use_linear=use_linear,
+                                use_checkpoint=use_checkpoint, only_self_att=temporal_self_att_only,
+                                causal_attention=use_causal_attention, relative_position=use_relative_position,
+                                temporal_length=temporal_length
+                            )
+                        )
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+                self.downsample_output.append(
+                    nn.Sequential(
+                        nn.GroupNorm(32, ch),
+                        nn.SiLU(),
+                        zero_module(conv_nd(dims, ch, ch, 3, padding=1))
+                    )
+                )
+            if level < len(channel_mult) - 1:
+                out_ch = ch
+                self.input_blocks.append(
+                    TimestepEmbedSequential(
+                        ResBlock(ch, time_embed_dim, dropout,
+                            out_channels=out_ch, dims=dims, use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            down=True
+                        )
+                        if resblock_updown
+                        else Downsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+                    )
+                )
+                self.downsample_output.append(
+                    nn.Sequential(
+                        nn.GroupNorm(32, out_ch),
+                        nn.SiLU(),
+                        zero_module(conv_nd(dims, out_ch, out_ch, 3, padding=1))
+                    )
+                )
+                ch = out_ch
+                ds *= 2
+
+    def forward(
+        self,
+        noisy_latents,
+        timesteps,
+        context_text,
+        context_img=None,
+        fps=None,
+        condition=None,            # [b, t, c, h, w]
+    ):
+        if self.ignore_noisy_latents:
+            noisy_latents = torch.zeros_like(noisy_latents)
+
+        b, _, t, height, width = noisy_latents.shape
+        t_emb = self.time_proj(timesteps).type(noisy_latents.dtype)
+        emb = self.time_embed(t_emb)
+
+        ## repeat t times for context [(b t) 77 768] & time embedding
+        ## check if we use per-frame image conditioning
+        if context_img is not None: ## decompose context into text and image
+            context_text = context_text.repeat_interleave(repeats=t, dim=0)
+            context_img = rearrange(context_img, 'b (t l) c -> (b t) l c', t=t)
+            context = torch.cat([context_text, context_img], dim=1)
+        else:
+            context = context_text.repeat_interleave(repeats=t, dim=0)
+        emb = emb.repeat_interleave(repeats=t, dim=0)
+
+        ## always in shape (b n t) c h w, except for temporal layer
+        noisy_latents = rearrange(noisy_latents, 'b c t h w -> (b t) c h w')
+        condition = rearrange(condition, 'b t c h w -> (b t) c h w')
+
+        ## combine emb
+        if self.fps_condition:
+            if fps is None:
+                fps = torch.tensor(
+                    [self.default_fs] * b, dtype=torch.long, device=noisy_latents.device)
+            fps_emb = self.time_proj(fps).type(noisy_latents.dtype)
+
+            fps_embed = self.fps_embedding(fps_emb)
+            fps_embed = fps_embed.repeat_interleave(repeats=t, dim=0)
+            emb = emb + fps_embed
+
+        h = noisy_latents.type(self.dtype)
+        hs = []
+        for id, module in enumerate(self.input_blocks):
+            h = module(h, emb, context=context, batch_size=b)
+            if id == 0:
+                h = h + self.cond_embedding(condition)
+                if self.addition_attention:
+                    h = self.init_attn(h, emb, context=context, batch_size=b)
+            hs.append(h)
+
+        guidance_feature_list = []
+        for hidden, module in zip(hs, self.downsample_output):
+            h = module(hidden)
+            guidance_feature_list.append(h)
+
+        return guidance_feature_list
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, layer_encoder_additional_kwargs={}, **kwargs):
+        cache_dir = kwargs.pop("cache_dir", None)
+        force_download = kwargs.pop("force_download", False)
+        proxies = kwargs.pop("proxies", None)
+        local_files_only = kwargs.pop("local_files_only", None)
+        token = kwargs.pop("token", None)
+        revision = kwargs.pop("revision", None)
+        subfolder = kwargs.pop("subfolder", None)
+        variant = kwargs.pop("variant", None)
+        use_safetensors = kwargs.pop("use_safetensors", None)
+
+        allow_pickle = False
+        if use_safetensors is None:
+            use_safetensors = True
+            allow_pickle = True
+
+        # Load config if we don't provide a configuration
+        config_path = pretrained_model_name_or_path
+
+        user_agent = {
+            "diffusers": __version__,
+            "file_type": "model",
+            "framework": "pytorch",
+        }
+
+        # load config
+        config, unused_kwargs, commit_hash = cls.load_config(
+            config_path,
+            cache_dir=cache_dir,
+            return_unused_kwargs=True,
+            return_commit_hash=True,
+            force_download=force_download,
+            proxies=proxies,
+            local_files_only=local_files_only,
+            token=token,
+            revision=revision,
+            subfolder=subfolder,
+            user_agent=user_agent,
+            **kwargs,
+        )
+
+        for key, value in layer_encoder_additional_kwargs.items():
+            if isinstance(value, (ListConfig, DictConfig)):
+                config[key] = OmegaConf.to_container(value, resolve=True)
+            else:
+                config[key] = value
+
+        # load model
+        model_file = None
+        if use_safetensors:
+            try:
+                model_file = _get_model_file(
+                    pretrained_model_name_or_path,
+                    weights_name=_add_variant(SAFETENSORS_WEIGHTS_NAME, variant),
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    proxies=proxies,
+                    local_files_only=local_files_only,
+                    token=token,
+                    revision=revision,
+                    subfolder=subfolder,
+                    user_agent=user_agent,
+                    commit_hash=commit_hash,
+                )
+
+            except IOError as e:
+                logger.error(f"An error occurred while trying to fetch {pretrained_model_name_or_path}: {e}")
+                if not allow_pickle:
+                    raise
+                logger.warning(
+                    "Defaulting to unsafe serialization. Pass `allow_pickle=False` to raise an error instead."
+                )
+
+        if model_file is None:
+            model_file = _get_model_file(
+                pretrained_model_name_or_path,
+                weights_name=_add_variant(WEIGHTS_NAME, variant),
+                cache_dir=cache_dir,
+                force_download=force_download,
+                proxies=proxies,
+                local_files_only=local_files_only,
+                token=token,
+                revision=revision,
+                subfolder=subfolder,
+                user_agent=user_agent,
+                commit_hash=commit_hash,
+            )
+
+        model = cls.from_config(config, **unused_kwargs)
+        state_dict = load_state_dict(model_file, variant)
+
+        missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+        print(f"Controlnet loaded from {model_file} with {len(missing_keys)} missing keys and {len(unexpected_keys)} unexpected keys.")
+        return model

lvdm/models/layer_controlnet.py

@@ -0,0 +1,444 @@
+from typing import Any, Dict, List, Optional, Tuple, Union
+from einops import rearrange, repeat
+import numpy as np
+from functools import partial
+import torch
+from torch import nn
+from torch.nn import functional as F
+from .unet import TimestepEmbedSequential, ResBlock, Downsample, Upsample, TemporalConvBlock
+from ..basics import zero_module, conv_nd
+from ..modules.attention import SpatialTransformer, TemporalTransformer
+from ..common import checkpoint
+
+from diffusers import __version__
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.model_loading_utils import load_state_dict
+from diffusers.utils import (
+    SAFETENSORS_WEIGHTS_NAME,
+    WEIGHTS_NAME,
+    logging,
+    _get_model_file,
+    _add_variant
+)
+from omegaconf import ListConfig, DictConfig, OmegaConf
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class ControlNetConditioningEmbedding(nn.Module):
+    """
+    Quoting from https://arxiv.org/abs/2302.05543: "Stable Diffusion uses a pre-processing method similar to VQ-GAN
+    [11] to convert the entire dataset of 512 × 512 images into smaller 64 × 64 “latent images” for stabilized
+    training. This requires ControlNets to convert image-based conditions to 64 × 64 feature space to match the
+    convolution size. We use a tiny network E(·) of four convolution layers with 4 × 4 kernels and 2 × 2 strides
+    (activated by ReLU, channels are 16, 32, 64, 128, initialized with Gaussian weights, trained jointly with the full
+    model) to encode image-space conditions ... into feature maps ..."
+    """
+
+    def __init__(
+        self,
+        conditioning_embedding_channels: int,
+        conditioning_channels: int = 3,
+        block_out_channels: Tuple[int, ...] = (16, 32, 96, 256),
+    ):
+        super().__init__()
+
+        self.conv_in = nn.Conv2d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1)
+
+        self.blocks = nn.ModuleList([])
+
+        for i in range(len(block_out_channels) - 1):
+            channel_in = block_out_channels[i]
+            channel_out = block_out_channels[i + 1]
+            self.blocks.append(nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1))
+            self.blocks.append(nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1, stride=2))
+
+        self.conv_out = zero_module(
+            nn.Conv2d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1)
+        )
+
+    def forward(self, conditioning):
+        embedding = self.conv_in(conditioning)
+        embedding = F.silu(embedding)
+
+        for block in self.blocks:
+            embedding = block(embedding)
+            embedding = F.silu(embedding)
+
+        embedding = self.conv_out(embedding)
+
+        return embedding
+
+
+class LayerControlNet(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels,
+        model_channels,
+        out_channels,
+        num_res_blocks,
+        attention_resolutions,
+        dropout=0.0,
+        channel_mult=(1, 2, 4, 8),
+        conv_resample=True,
+        dims=2,
+        context_dim=None,
+        use_scale_shift_norm=False,
+        resblock_updown=False,
+        num_heads=-1,
+        num_head_channels=-1,
+        transformer_depth=1,
+        use_linear=False,
+        use_checkpoint=False,
+        temporal_conv=False,
+        tempspatial_aware=False,
+        temporal_attention=True,
+        use_relative_position=True,
+        use_causal_attention=False,
+        temporal_length=None,
+        addition_attention=False,
+        temporal_selfatt_only=True,
+        image_cross_attention=False,
+        image_cross_attention_scale_learnable=False,
+        default_fps=4,
+        fps_condition=False,
+        ignore_noisy_latents=True,
+        condition_channels={},
+        control_injection_mode='add',
+        use_vae_for_trajectory=False,
+    ):
+        super().__init__()
+        if num_heads == -1:
+            assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
+        if num_head_channels == -1:
+            assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
+
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.temporal_attention = temporal_attention
+        time_embed_dim = model_channels * 4
+        self.use_checkpoint = use_checkpoint
+        temporal_self_att_only = True
+        self.addition_attention = addition_attention
+        self.temporal_length = temporal_length
+        self.image_cross_attention = image_cross_attention
+        self.image_cross_attention_scale_learnable = image_cross_attention_scale_learnable
+        self.default_fps = default_fps
+        self.fps_condition = fps_condition
+        self.ignore_noisy_latents = ignore_noisy_latents
+        assert len(condition_channels) > 0, 'Condition types must be specified'
+        self.condition_channels = condition_channels
+        self.control_injection_mode = control_injection_mode
+        self.use_vae_for_trajectory = use_vae_for_trajectory
+
+        ## Time embedding blocks
+        self.time_proj = Timesteps(model_channels, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.time_embed = TimestepEmbedding(model_channels, time_embed_dim)
+
+        if fps_condition:
+            self.fps_embedding = TimestepEmbedding(model_channels, time_embed_dim)
+            nn.init.zeros_(self.fps_embedding.linear_2.weight)
+            nn.init.zeros_(self.fps_embedding.linear_2.bias)
+
+        if "motion_score" in condition_channels:
+            if control_injection_mode == 'add':
+                self.motion_embedding = zero_module(conv_nd(dims, condition_channels["motion_score"], model_channels, 3, padding=1))
+            elif control_injection_mode == 'concat':
+                self.motion_embedding = zero_module(conv_nd(dims, condition_channels["motion_score"], condition_channels["motion_score"], 3, padding=1))
+            else:
+                raise ValueError(f"control_injection_mode {control_injection_mode} is not supported, use 'add' or 'concat'")
+        if "sketch" in condition_channels:
+            if control_injection_mode == 'add':
+                self.sketch_embedding = zero_module(conv_nd(dims, condition_channels["sketch"], model_channels, 3, padding=1))
+            elif control_injection_mode == 'concat':
+                self.sketch_embedding = zero_module(conv_nd(dims, condition_channels["sketch"], condition_channels["sketch"], 3, padding=1))
+            else:
+                raise ValueError(f"control_injection_mode {control_injection_mode} is not supported, use 'add' or 'concat'")
+        if "trajectory" in condition_channels:
+            if control_injection_mode == 'add':
+                if use_vae_for_trajectory:
+                    self.trajectory_embedding = zero_module(conv_nd(dims, condition_channels["trajectory"], model_channels, 3, padding=1))
+                else:
+                    self.trajectory_embedding = ControlNetConditioningEmbedding(model_channels, condition_channels["trajectory"])
+            elif control_injection_mode == 'concat':
+                if use_vae_for_trajectory:
+                    self.trajectory_embedding = zero_module(conv_nd(dims, condition_channels["trajectory"], condition_channels["trajectory"], 3, padding=1))
+                else:
+                    self.trajectory_embedding = ControlNetConditioningEmbedding(condition_channels["trajectory"], condition_channels["trajectory"])
+            else:
+                raise ValueError(f"control_injection_mode {control_injection_mode} is not supported, use 'add' or 'concat'")
+
+        self.input_blocks = nn.ModuleList(
+            [
+                TimestepEmbedSequential(conv_nd(dims, in_channels, model_channels, 3, padding=1))
+            ]
+        )
+
+        if self.addition_attention:
+            self.init_attn = TimestepEmbedSequential(
+                TemporalTransformer(
+                    model_channels,
+                    n_heads=8,
+                    d_head=num_head_channels,
+                    depth=transformer_depth,
+                    context_dim=context_dim,
+                    use_checkpoint=use_checkpoint, only_self_att=temporal_selfatt_only,
+                    causal_attention=False, relative_position=use_relative_position,
+                    temporal_length=temporal_length
+                )
+            )
+
+        ch = model_channels
+        ds = 1
+        for level, mult in enumerate(channel_mult):
+            for _ in range(num_res_blocks):
+                layers = [
+                    ResBlock(ch, time_embed_dim, dropout,
+                        out_channels=mult * model_channels, dims=dims, use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm, tempspatial_aware=tempspatial_aware,
+                        use_temporal_conv=temporal_conv
+                    )
+                ]
+                ch = mult * model_channels
+                if ds in attention_resolutions:
+                    if num_head_channels == -1:
+                        dim_head = ch // num_heads
+                    else:
+                        num_heads = ch // num_head_channels
+                        dim_head = num_head_channels
+                    layers.append(
+                        SpatialTransformer(ch, num_heads, dim_head,
+                            depth=transformer_depth, context_dim=context_dim, use_linear=use_linear,
+                            use_checkpoint=use_checkpoint, disable_self_attn=False,
+                            video_length=temporal_length, image_cross_attention=self.image_cross_attention,
+                            image_cross_attention_scale_learnable=self.image_cross_attention_scale_learnable,
+                        )
+                    )
+                    if self.temporal_attention:
+                        layers.append(
+                            TemporalTransformer(ch, num_heads, dim_head,
+                                depth=transformer_depth, context_dim=context_dim, use_linear=use_linear,
+                                use_checkpoint=use_checkpoint, only_self_att=temporal_self_att_only,
+                                causal_attention=use_causal_attention, relative_position=use_relative_position,
+                                temporal_length=temporal_length
+                            )
+                        )
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+
+            if level < len(channel_mult) - 1:
+                out_ch = ch
+                self.input_blocks.append(
+                    TimestepEmbedSequential(
+                        ResBlock(ch, time_embed_dim, dropout,
+                            out_channels=out_ch, dims=dims, use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            down=True
+                        )
+                        if resblock_updown
+                        else Downsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+                    )
+                )
+                ch = out_ch
+                ds *= 2
+
+    def forward(
+        self,
+        noisy_latents,
+        timesteps,
+        context_text,
+        context_img=None,
+        fps=None,
+        layer_latents=None,     # [b, n_layer, t, c, h, w]
+        layer_latent_mask=None, # [b, n_layer, t, 1, h, w]
+        motion_scores=None,     # [b, n_layer]
+        sketch=None,            # [b, n_layer, t, c, h, w]
+        trajectory=None,        # [b, n_layer, t, c, h, w]
+    ):
+        if self.ignore_noisy_latents:
+            noisy_latents_shape = list(noisy_latents.shape)
+            noisy_latents_shape[1] = 0
+            noisy_latents = torch.zeros(noisy_latents_shape, device=noisy_latents.device, dtype=noisy_latents.dtype)
+
+        b, _, t, height, width = noisy_latents.shape
+        n_layer = layer_latents.shape[1]
+        t_emb = self.time_proj(timesteps).type(noisy_latents.dtype)
+        emb = self.time_embed(t_emb)
+
+        ## repeat t times for context [(b t) 77 768] & time embedding
+        ## check if we use per-frame image conditioning
+        if context_img is not None: ## decompose context into text and image
+            context_text = repeat(context_text, 'b l c -> (b n t) l c', n=n_layer, t=t)
+            context_img = repeat(context_img, 'b tl c -> b n tl c', n=n_layer)
+            context_img = rearrange(context_img, 'b n (t l) c -> (b n t) l c', t=t)
+            context = torch.cat([context_text, context_img], dim=1)
+        else:
+            context = repeat(context_text, 'b l c -> (b n t) l c', n=n_layer, t=t)
+        emb = repeat(emb, 'b c -> (b n t) c', n=n_layer, t=t)
+
+        ## always in shape (b n t) c h w, except for temporal layer
+        noisy_latents = repeat(noisy_latents, 'b c t h w -> (b n t) c h w', n=n_layer)
+
+        ## combine emb
+        if self.fps_condition:
+            if fps is None:
+                fps = torch.tensor(
+                    [self.default_fs] * b, dtype=torch.long, device=noisy_latents.device)
+            fps_emb = self.time_proj(fps).type(noisy_latents.dtype)
+
+            fps_embed = self.fps_embedding(fps_emb)
+            fps_embed = repeat(fps_embed, 'b c -> (b n t) c', n=n_layer, t=t)
+            emb = emb + fps_embed
+
+        ## process conditions
+        layer_condition = torch.cat([layer_latents, layer_latent_mask], dim=3)
+        layer_condition = rearrange(layer_condition, 'b n t c h w -> (b n t) c h w')
+        h = torch.cat([noisy_latents, layer_condition], dim=1)
+
+        if "motion_score" in self.condition_channels:
+            motion_condition = repeat(motion_scores, 'b n -> b n t 1 h w', t=t, h=height, w=width)
+            motion_condition = torch.cat([motion_condition, layer_latent_mask], dim=3)
+            motion_condition = rearrange(motion_condition, 'b n t c h w -> (b n t) c h w')
+            motion_condition = self.motion_embedding(motion_condition)
+            if self.control_injection_mode == 'concat':
+                h = torch.cat([h, motion_condition], dim=1)
+
+        if "sketch" in self.condition_channels:
+            sketch_condition = rearrange(sketch, 'b n t c h w -> (b n t) c h w')
+            sketch_condition = self.sketch_embedding(sketch_condition)
+            if self.control_injection_mode == 'concat':
+                h = torch.cat([h, sketch_condition], dim=1)
+
+        if "trajectory" in self.condition_channels:
+            traj_condition = rearrange(trajectory, 'b n t c h w -> (b n t) c h w')
+            traj_condition = self.trajectory_embedding(traj_condition)
+            if self.control_injection_mode == 'concat':
+                h = torch.cat([h, traj_condition], dim=1)
+
+        layer_features = []
+        for id, module in enumerate(self.input_blocks):
+            h = module(h, emb, context=context, batch_size=b*n_layer)
+            if id == 0:
+                if self.control_injection_mode == 'add':
+                    if "motion_score" in self.condition_channels:
+                        h = h + motion_condition
+                    if "sketch" in self.condition_channels:
+                        h = h + sketch_condition
+                    if "trajectory" in self.condition_channels:
+                        h = h + traj_condition
+                if self.addition_attention:
+                    h = self.init_attn(h, emb, context=context, batch_size=b*n_layer)
+            if SpatialTransformer in [type(m) for m in module]:
+                layer_features.append(rearrange(h, '(b n t) c h w -> b n t c h w', b=b, n=n_layer))
+
+        return layer_features
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, layer_controlnet_additional_kwargs={}, **kwargs):
+        cache_dir = kwargs.pop("cache_dir", None)
+        force_download = kwargs.pop("force_download", False)
+        proxies = kwargs.pop("proxies", None)
+        local_files_only = kwargs.pop("local_files_only", None)
+        token = kwargs.pop("token", None)
+        revision = kwargs.pop("revision", None)
+        subfolder = kwargs.pop("subfolder", None)
+        variant = kwargs.pop("variant", None)
+        use_safetensors = kwargs.pop("use_safetensors", None)
+
+        allow_pickle = False
+        if use_safetensors is None:
+            use_safetensors = True
+            allow_pickle = True
+
+        # Load config if we don't provide a configuration
+        config_path = pretrained_model_name_or_path
+
+        user_agent = {
+            "diffusers": __version__,
+            "file_type": "model",
+            "framework": "pytorch",
+        }
+
+        # load config
+        config, unused_kwargs, commit_hash = cls.load_config(
+            config_path,
+            cache_dir=cache_dir,
+            return_unused_kwargs=True,
+            return_commit_hash=True,
+            force_download=force_download,
+            proxies=proxies,
+            local_files_only=local_files_only,
+            token=token,
+            revision=revision,
+            subfolder=subfolder,
+            user_agent=user_agent,
+            **kwargs,
+        )
+
+        for key, value in layer_controlnet_additional_kwargs.items():
+            if isinstance(value, (ListConfig, DictConfig)):
+                config[key] = OmegaConf.to_container(value, resolve=True)
+            else:
+                config[key] = value
+
+        # load model
+        model_file = None
+        if use_safetensors:
+            try:
+                model_file = _get_model_file(
+                    pretrained_model_name_or_path,
+                    weights_name=_add_variant(SAFETENSORS_WEIGHTS_NAME, variant),
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    proxies=proxies,
+                    local_files_only=local_files_only,
+                    token=token,
+                    revision=revision,
+                    subfolder=subfolder,
+                    user_agent=user_agent,
+                    commit_hash=commit_hash,
+                )
+
+            except IOError as e:
+                logger.error(f"An error occurred while trying to fetch {pretrained_model_name_or_path}: {e}")
+                if not allow_pickle:
+                    raise
+                logger.warning(
+                    "Defaulting to unsafe serialization. Pass `allow_pickle=False` to raise an error instead."
+                )
+
+        if model_file is None:
+            model_file = _get_model_file(
+                pretrained_model_name_or_path,
+                weights_name=_add_variant(WEIGHTS_NAME, variant),
+                cache_dir=cache_dir,
+                force_download=force_download,
+                proxies=proxies,
+                local_files_only=local_files_only,
+                token=token,
+                revision=revision,
+                subfolder=subfolder,
+                user_agent=user_agent,
+                commit_hash=commit_hash,
+            )
+
+        model = cls.from_config(config, **unused_kwargs)
+        state_dict = load_state_dict(model_file, variant)
+
+        if state_dict['input_blocks.0.0.weight'].shape[1] != model.input_blocks[0][0].weight.shape[1]:
+            state_dict.pop('input_blocks.0.0.weight')
+
+        missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+        print(f"LayerControlNet loaded from {model_file} with {len(missing_keys)} missing keys and {len(unexpected_keys)} unexpected keys.")
+        return model