add 1333 script

ckpt
Resolve OOM
2026-03-02 21:50:31 +00:00 · 2026-02-09 22:41:25 +00:00 · 2025-11-27 21:08:32 +00:00 · 2025-11-25 09:33:33 +00:00 · 2025-10-13 21:10:34 +00:00 · 2025-10-11 09:58:20 +00:00
97 changed files with 3176 additions and 465 deletions
@@ -0,0 +1,452 @@
+{
+  "data": [
+    {
+      "prompt": "Young man skating with a skateboard on the ramps with graffiti of a park with trees, on a sunny day.", 
+      "image_path": "images/mixkit-boy-skating-with-a-skateboard-in-a-park-with-ramps-34389.png"
+    },
+    {
+      "prompt": "In the midst of the joyous New Year's Eve celebration, the cheerful group of friends, their spirits lifted by the festivities, decides to immortalize the moment with a vibrant snapshot", 
+      "image_path": "images/mixkit-a-cheerful-group-of-friends-celebrate-new-years-eve-and-51525.png"
+    },
+    {
+      "prompt": "A man and a woman playing in a field with grass, during a bright afternoon, while cars pass by in the distance.", 
+      "image_path": "images/mixkit-a-cute-couple-playing-on-the-grass-4688.png"
+    },
+    {
+      "prompt": "Aerial view of a rocky mountain in the forest at a sunny day drone flight footage", 
+      "image_path": "images/mixkit-aerial-view-of-a-rocky-mountain-in-the-forest-50589.png"
+    },
+    {
+      "prompt": "A little girl wearing a pink security helmet and denim overall discovers the art of cycling amidst the serene park, as the camera captures her graceful progress.", 
+      "image_path": "images/mixkit-a-little-girl-cruises-through-the-forest-path-on-her-50088.png"
+    },
+    {
+      "prompt": "Aerial shot of a beach shore with sea waves. Big rocks on the sand at an alone beach.", 
+      "image_path": "images/mixkit-aerial-shot-of-a-beach-with-sea-waves-1087.png"
+    },
+    {
+      "prompt": "Young woman cleaning her house decorated with plants and decorations, while dancing happily to music in her headphones.", 
+      "image_path": "images/mixkit-woman-cleaning-her-house-dancing-happy-43379.png"
+    },
+    {
+      "prompt": "Aerial tour in a meadow surrounded by hills on the horizon, while some birds fly low over a lake.", 
+      "image_path": "images/mixkit-birds-flying-low-over-a-lake-in-a-meadow-41417.png"
+    },
+    {
+      "prompt": "In the video, a lone rider guides a majestic horse across an expansive, open field as the sun sets in the background. The rider, dressed in a classic blue shirt and wide-brimmed hat, sits confidently in the saddle, silhouetted against the warm glow of the evening sky. The horse moves gracefully, its mane and tail flowing with each step, creating a sense of harmony between horse and rider. Surrounding the pair, towering trees form a natural border, their leaves gently rustling in the breeze. The shadows lengthen on the ground, accentuating the serene and timeless feel of the scene. The distant hills and wooden fences frame the horizon, adding depth to the tranquil landscape. A few horses graze peacefully in the background, blending into the pastoral setting. The overall ambiance evokes a sense of calmness and quietude, capturing a perfect moment in the golden light of dusk.", 
+      "image_path": "images/mixkit-a-rancher-riding-a-horse-at-sunset-1143.png"
+    },
+    {
+      "prompt": "In the video, a martial artist dressed in a traditional white uniform with a black belt demonstrates a series of precise movements against a stark black background. The individual gracefully transitions between stances, embodying a sense of focused discipline and control. Each motion is executed with a deliberate pace, showcasing the fluidity of martial arts techniques. The soft lighting creates subtle highlights on the uniform, adding depth to the figure as it moves. The practitioner begins with an open-hand pose, feet firmly grounded, gradually shifting to a powerful forward punch. The fluidity of the sequence displays a mastery of balance and poise. Every trajectory of the limbs is precise and deliberate, capturing the elegance and strength of martial arts. The serene, isolated setting enhances the intensity and concentration of the practitioner. This visual presentation is an elegant interplay of motion and stillness, displaying the art form's discipline and grace.", 
+      "image_path": "images/mixkit-a-young-man-practicing-his-karate-moves-49635.png"
+    },
+    {
+      "prompt": "In a serene and softly lit yoga studio, three individuals engage in a yoga session, each performing an upward-facing stretch. The central figure is a woman with shoulder-length brown hair, dressed in a light cropped top and green leggings, her posture reflecting grace and concentration. To her right, another participant, a woman in a purple outfit, mirrors the pose with equal poise. On her left, a person with a bun focuses intently, supported slightly by yoga blocks beneath their hands. The warm-colored wooden floor contrasts soothingly with the soft pastel mural on the back wall, featuring an abstract design and partial visage of a serene face. Natural light floods the space from a large window on the right, where lush greens peek through, adding an element of tranquility. In the corner of the room, a collection of meditation instruments, including a gong and a Buddha statue, subtly frame the peaceful setting. The mood is calm yet focused, as all three participants are deeply engaged in their practice. The scene combines elements of balance, harmony, and a shared journey towards mindfulness. This depiction captures the essence of a yoga session that blends personal growth with collective experience.", 
+      "image_path": "images/mixkit-small-group-of-people-doing-yoga-together-43730.png"
+    },
+    {
+      "prompt": "In the deep blue expanse of the ocean, two dolphins glide effortlessly, their sleek bodies reflecting the sunlight filtering through the water. The prominent shadows and caustics create a shimmering effect on their skin, capturing the beauty of their natural habitat. Each dolphin moves with a fluid grace, occasionally interacting with gentle nudges, showcasing their playful and social nature. The scene is vibrant and dynamic, with the clear blue background accentuating the dolphins' movements, making it an ideal subject for AI recreation.", 
+      "image_path": "images/mixkit-dolphins-underwater-4133.png"
+    },
+    {
+      "prompt": "A bustling ski slope comes alive with skiers descending a pristine, snow-covered hill, surrounded by towering, snow-draped evergreens. Several figures stand atop the slope, silhouetted against a clear blue sky, preparing to embark on their ski run. The chair lift on the right continuously drops off eager adventurers, adding to the excitement at the hilltop. Each skier, clad in colorful winter gear, carves distinct paths into the textured snow as they weave their way down. The interplay of sunlight and shadows accentuates the myriad tracks etched into the slope, creating a dynamic visual rhythm. The scene captures a vibrant winter wonderland, full of action and the thrill of a perfect ski day.", 
+      "image_path": "images/mixkit-skiers-on-a-snowy-slope-3327.png"
+    },
+    {
+      "prompt": "A determined climber is scaling a massive rock face, showcasing exceptional strength and skill. The person, clad in a teal shirt and dark pants, climbs with precision, their movements measured and deliberate. They are secured by climbing gear, which includes ropes and a harness, emphasizing their commitment to safety. The rugged texture of the sandy-colored rock provides an imposing backdrop, adding drama and scale to the climb. In the distance, other large rock formations and sparse vegetation can be seen under a bright, overcast sky, contributing to the natural and adventurous atmosphere. The scene captures a moment of focus and challenge, highlighting the climber's tenacity and the breathtaking environment.",
+      "image_path": "images/mixkit-alpinist-climbing-a-huge-rock-in-a-desert-43306.png"
+    },
+    {
+      "prompt": "A silver SUV drives along a winding, snow-covered mountain road, with dense pine trees blanketed in snow lining both sides. The scene is serene, with the vehicle moving smoothly, possibly on a winter journey or vacation. As the SUV disappears around the bend, another, darker SUV follows, creating a sense of motion and perspective on the snow-dusted asphalt. The towering, snow-laden rock formation to the right contrasts with the dark green of the pines, highlighting the peacefulness of the wintry landscape.",
+      "image_path": "images/mixkit-curve-on-a-snowy-forest-road-3317.png"
+    },
+    {
+      "prompt": "A solitary boat glides across the expansive, tranquil expanse of a serene lake. The vessel leaves a gentle wake behind, creating delicate ripples across the mirror-like surface. The water appears a rich shade of teal, seamlessly blending with the sky at the horizon. Silhouettes of distant trees are faintly visible, creating a picturesque backdrop that enhances the solitary journey of the boat. The sky is a calm gradient, shifting from soft oranges near the shore to the pale blues above. In the distance, a few slender poles emerge from the water, remnants of an old structure or natural formation. The mood of the scene is one of peace and solitude, with the boat journeying steadily through the quiet landscape. There is a sense of endless possibilities as the boat moves toward the unseen beyond the frame. The simplicity and stillness of the scene invite contemplation and reflection, encapsulating a perfect moment of quietude on the water.",
+      "image_path": "images/mixkit-motorboat-on-a-large-lake-with-turquoise-blue-waters-4996.png"
+    },
+    {
+      "prompt": "Far shot from high behind a beach on a sunny day, with crowds of people everywhere, buildings and hotels behind, and the sea on the horizon as ships cross it.",
+      "image_path": "480P_i2v_mixkit_images/Far shot from high behind a beach on a sunny day, with crowds of people everywhere, buildings and ho.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Skillful cyclist doing a wheelie on a bike while riding through a forest, on a dirt road, surrounded by many trees, in the morning.",
+      "image_path": "480P_i2v_mixkit_images/Skillful cyclist doing a wheelie on a bike while riding through a forest, on a dirt road, surrounded.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Portrait of a man with a mask and fluorescent paint dancing in front of the camera, under a blue party light.",
+      "image_path": "480P_i2v_mixkit_images/Portrait of a man with a mask and fluorescent paint dancing in front of the camera, under a blue par.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Timelapse of a luxury swimming pool in a boutique hotel.",
+      "image_path": "480P_i2v_mixkit_images/Timelapse of a luxury swimming pool in a boutique hotel..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Man dressed in 80's style dances very happily in his kitchen while listening to music on his radio and drinking wine.",
+      "image_path": "480P_i2v_mixkit_images/Man dressed in 80's style dances very happily in his kitchen while listening to music on his radio a.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Romantic couple embracing and looking at each other in the middle of a forest, during a break on a road trip through nature.",
+      "image_path": "480P_i2v_mixkit_images/Romantic couple embracing and looking at each other in the middle of a forest, during a break on a r.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Big city with many big buildings and a big capitol, in a view from a slow drive through the heights in the afternoon.",
+      "image_path": "480P_i2v_mixkit_images/Big city with many big buildings and a big capitol, in a view from a slow drive through the heights .jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "A man appears in different parts of an abandoned building, standing in the dark, swinging a baseball bat, outside.",
+      "image_path": "480P_i2v_mixkit_images/A man appears in different parts of an abandoned building, standing in the dark, swinging a baseball.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Many small lights out of focus yellow and blue and of different intensities, at night on a blue background.",
+      "image_path": "480P_i2v_mixkit_images/Many small lights out of focus yellow and blue and of different intensities, at night on a blue back.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Pair of jazz musicians performing a song with their saxophone and trombone on an abandoned train.",
+      "image_path": "480P_i2v_mixkit_images/Pair of jazz musicians performing a song with their saxophone and trombone on an abandoned train..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Pair of mountain bikers climbing a slope in the forest covered with rocks, with abundant trees and vegetation around.",
+      "image_path": "480P_i2v_mixkit_images/Pair of mountain bikers climbing a slope in the forest covered with rocks, with abundant trees and v.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "A model with cyberpunk makeup poses in front of a neon light circle background.",
+      "image_path": "480P_i2v_mixkit_images/A model with cyberpunk makeup poses in front of a neon light circle background..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Couple of boy and girl walking hand in hand through a meadow among many cows that feed on grass and grass.",
+      "image_path": "480P_i2v_mixkit_images/Couple of boy and girl walking hand in hand through a meadow among many cows that feed on grass and .jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Natural coastal relief on a sunny day, with large hills covered with dry grass in an arid climate, in a panoramic aerial view.",
+      "image_path": "480P_i2v_mixkit_images/Natural coastal relief on a sunny day, with large hills covered with dry grass in an arid climate, i.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Young woman lying on her bed, while chatting on her cell phone, in a close-up shot of her hands typing on the phone.",
+      "image_path": "480P_i2v_mixkit_images/Young woman lying on her bed, while chatting on her cell phone, in a close-up shot of her hands typi.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Meadow with a frozen lake in a forest of pines covered in snow and fog.",
+      "image_path": "480P_i2v_mixkit_images/Meadow with a frozen lake in a forest of pines covered in snow and fog..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Aerial shot of buses and cars on the freeway with a cloudy sky in the background.",
+      "image_path": "480P_i2v_mixkit_images/Aerial shot of buses and cars on the freeway with a cloudy sky in the background..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Beautiful countryside meadow during summer with trees, grass and little yellow flowers in the ground on a sunny day.",
+      "image_path": "480P_i2v_mixkit_images/Beautiful countryside meadow during summer with trees, grass and little yellow flowers in the ground.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "A man wearing grey shorts jumps rope in a gym, weights and gym equipment in the background.",
+      "image_path": "480P_i2v_mixkit_images/A man wearing grey shorts jumps rope in a gym, weights and gym equipment in the background..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "A young woman wearing a black mini skirt, a red sleeveless shirt and a red Santa hat stands on a stage and sings into a microphone. In the background, a guitar player.",
+      "image_path": "480P_i2v_mixkit_images/A young woman wearing a black mini skirt, a red sleeveless shirt and a red Santa hat stands on a sta.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Little baby with a pacifier, playing with a teddy bear, accompanied by his mom, both sitting on a bed, in a close view.",
+      "image_path": "480P_i2v_mixkit_images/Little baby with a pacifier, playing with a teddy bear, accompanied by his mom, both sitting on a be.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "A saxophonist wearing a blazer dances while playing a song in a park.",
+      "image_path": "480P_i2v_mixkit_images/A saxophonist wearing a blazer dances while playing a song in a park..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "A young woman with short hair wearing pink sunglasses chews gum and makes a bubble gum with the city in the background.",
+      "image_path": "480P_i2v_mixkit_images/A young woman with short hair wearing pink sunglasses chews gum and makes a bubble gum with the city.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "White flowers from a bridal bouquet, with the background out of focus.",
+      "image_path": "480P_i2v_mixkit_images/White flowers from a bridal bouquet, with the background out of focus..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Silhouette of a person looking at an image projected in front, with spectra and light patterns rotating around a center in the background.",
+      "image_path": "480P_i2v_mixkit_images/Silhouette of a person looking at an image projected in front, with spectra and light patterns rotat.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Natural aerial landscape with a relief covered with abundant trees and vegetation and a thick layer of mist.",
+      "image_path": "480P_i2v_mixkit_images/Natural aerial landscape with a relief covered with abundant trees and vegetation and a thick layer .jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Boy dancing happily and cheerfully in a room set in the 80s.",
+      "image_path": "480P_i2v_mixkit_images/Boy dancing happily and cheerfully in a room set in the 80s..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Close up video of the legs of a young gymnast woman doing some acrobatics on a balance beam in a dark and empty gym.",
+      "image_path": "480P_i2v_mixkit_images/Close up video of the legs of a young gymnast woman doing some acrobatics on a balance beam in a dar.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Female scientist in a laboratory, with test tube, and more containers, looking through the lens of a microscope.",
+      "image_path": "480P_i2v_mixkit_images/Female scientist in a laboratory, with test tube, and more containers, looking through the lens of a.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Young woman getting makeup done by a professional makeup artist for a winter concept fashion photo shoot.",
+      "image_path": "480P_i2v_mixkit_images/Young woman getting makeup done by a professional makeup artist for a winter concept fashion photo s.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Loving couple sitting on a log on the shore of a lake outside, sharing an affectionate hug.",
+      "image_path": "480P_i2v_mixkit_images/Loving couple sitting on a log on the shore of a lake outside, sharing an affectionate hug..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Small family of father, mother and daughter, walking together through nature in a forest.",
+      "image_path": "480P_i2v_mixkit_images/Small family of father, mother and daughter, walking together through nature in a forest..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Flying over a peninsula covered in bushy trees, while discovering the sea around it, painted a beautiful turquoise blue, on a sunny day.",
+      "image_path": "480P_i2v_mixkit_images/Flying over a peninsula covered in bushy trees, while discovering the sea around it, painted a beaut.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "A woman with long blonde hair wearing a white t-shirt and black bikini bottoms stands on a wooden deck smiling, with the beach and palm trees in the background.",
+      "image_path": "480P_i2v_mixkit_images/A woman with long blonde hair wearing a white t-shirt and black bikini bottoms stands on a wooden de.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "View towards the horizon in a flat desert at night, as the full moon glows brightly red.",
+      "image_path": "480P_i2v_mixkit_images/View towards the horizon in a flat desert at night, as the full moon glows brightly red..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "100 dollar bill thrown in the street seen up close with the background out of focus, while a person arrives, picks it up and walks away.",
+      "image_path": "480P_i2v_mixkit_images/100 dollar bill thrown in the street seen up close with the background out of focus, while a person .jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Flying over a forest of abundant trees and vegetation, with a river and a dam, with houses and hills in the surroundings.",
+      "image_path": "480P_i2v_mixkit_images/Flying over a forest of abundant trees and vegetation, with a river and a dam, with houses and hills.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Artistic portrait of a man with his body and face stained with phosphorescent paint that glows with black light in a dark place.",
+      "image_path": "480P_i2v_mixkit_images/Artistic portrait of a man with his body and face stained with phosphorescent paint that glows with .jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Aerial view of the calm blue sea made with a drone showing the horizon and the sea.",
+      "image_path": "480P_i2v_mixkit_images/Aerial view of the calm blue sea made with a drone showing the horizon and the sea..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Cars driving on a two way street underneath a tree, with traffic signs, buildings and the sun in the background.",
+      "image_path": "480P_i2v_mixkit_images/Cars driving on a two way street underneath a tree, with traffic signs, buildings and the sun in the.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Face of a young woman, close up portrait",
+      "image_path": "480P_i2v_mixkit_images/Face of a young woman, close up portrait.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Man with his body splattered with fluorescent paint while doing a slow conceptual dance, under a black light in a dark place with colored spotlights around.",
+      "image_path": "480P_i2v_mixkit_images/Man with his body splattered with fluorescent paint while doing a slow conceptual dance, under a bla.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Two women dressed in colorful retro styles, having fun taking photos in a room with lots of pastel colors.",
+      "image_path": "480P_i2v_mixkit_images/Two women dressed in colorful retro styles, having fun taking photos in a room with lots of pastel c.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "The silhouette of a person who moves his head reflects inside a city with tall buildings next to a lake.",
+      "image_path": "480P_i2v_mixkit_images/The silhouette of a person who moves his head reflects inside a city with tall buildings next to a l.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "View from a car window of the trees lining the road speeding by, with the dazzling sunlight in the sky, in a retro looking video.",
+      "image_path": "480P_i2v_mixkit_images/View from a car window of the trees lining the road speeding by, with the dazzling sunlight in the s.jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Aerial view of two people on a tennis court playing tennis.",
+      "image_path": "480P_i2v_mixkit_images/Aerial view of two people on a tennis court playing tennis..jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "View from the floor upwards of a young woman with long dark hair dancing slowly in green and purple neon lights",
+      "image_path": "480P_i2v_mixkit_images/View from the floor upwards of a young woman with long dark hair dancing slowly in green and purple .jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "prompt": "Some friends dancing and having fun together in circles, at a party surrounded by colored lights at a party, in a fancy old place, in a view from below them.",
+      "image_path": "480P_i2v_mixkit_images/Some friends dancing and having fun together in circles, at a party surrounded by colored lights at .jpg",
+      "num_inference_steps": 3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    }
+  ]
+}
@@ -0,0 +1,185 @@
+#!/bin/bash
+#SBATCH --job-name=1.3B_high_noise_moe_8n_sf_distill
+#SBATCH --partition=main
+#SBATCH --nodes=8
+#SBATCH --ntasks=8
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=1.3B_moe_8n_sf_distill_output/moe_sf_distill_%j.out
+#SBATCH --error=1.3B_moe_8n_sf_distill_output/moe_sf_distill_%j.err
+#SBATCH --exclusive
+set -e -x
+
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate wei-fv
+export HOME="/mnt/weka/home/hao.zhang/wei"
+
+# Basic Info
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29500
+export NODE_RANK=$SLURM_PROCID
+nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
+export MASTER_ADDR=${nodes[0]}
+export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
+export TOKENIZERS_PARALLELISM=false
+export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+# Configs
+NUM_GPUS=8
+NUM_TOTAL_GPUS=64
+
+# Model paths for Self-Forcing DMD distillation:
+GENERATOR_MODEL_PATH="rand0nmr/SFWan2.1-T2V-A1.3B-Diffusers"
+# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"  # Teacher model
+REAL_SCORE_MODEL_PATH="/mnt/sharefs/users/hao.zhang/wei/wan_models/Wan2.1-T2V-14B-Diffusers"
+FAKE_SCORE_MODEL_PATH="/mnt/sharefs/users/hao.zhang/wei/wan_models/Wan2.1-T2V-A1.3B-Diffusers"
+
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+# DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
+DATA_DIR="/mnt/sharefs/users/hao.zhang/wl/release/vidprom_16k_text_embed"
+# DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+# VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_2.json"
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  # Updated for self-forcing DMD
+  --output_dir "/mnt/sharefs/users/hao.zhang/wei/1.3B_moe_x0_CacheAppend_3b_4c_high_x0_ode_init_addnoisehigh_moecritic"
+  --wandb_run_name "1.3B_moe_8n_sf_distill_x0_CacheAppend_3b_4c_high_x0_ode_init_addnoisehigh"
+  # --use_sf_wan
+  # --sf_ode_init_path "checkpoints/ode_init.pt"
+  --max_train_steps 5000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 81  # Must be divisible by num_frame_per_block (81 % 3 = 0 ✓)
+  --enable_gradient_checkpointing_type "full"
+  --log_visualization
+  --simulate_generator_forward
+  --num_frame_per_block 3  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 21
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_TOTAL_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
+)
+
+# Model arguments
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH  # TODO: check if you can remove this in this script
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --generator_model_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 8
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 100
+  --validation_sampling_steps "4"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+  --vae_cpu_offload True
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-6
+  --fake_score_learning_rate 4e-7
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 200
+  --weight_only_checkpointing_steps 200
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 10.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 5
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 200
+  # --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wei/FastVideo/diffusers_1.3B_SF/model.safetensors"
+  # --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.saftensors"
+  --init_weights_from_safetensors "/mnt/sharefs/users/hao.zhang/wl/release/official_ode_init_8n_high_1_3_x0/checkpoint_model_002000/model.safetensors"
+  --init_weights_from_safetensors_2 "/mnt/sharefs/users/hao.zhang/wl/release/official_ode_init_8n_low_1_3/checkpoint_model_002000/model.safetensors"
+  # --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/FastVideo2/vidprom_8b16k_1e-5_gn1/checkpoint-3500/transformer/diffusion_pytorch_model.safetensors"
+#   --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/sf2/checkpoint_v3_7k/model.safetensors"
+  # --resume_from_checkpoint "/mnt/sharefs/users/hao.zhang/wei/1.3B_MoE_SFwan_t2v_finetune/checkpoint-1990"
+)
+
+# Self-forcing DMD arguments
+dmd_args=(
+  --dmd_denoising_steps '1000,750,500,250'
+  --min_timestep_ratio 0.04
+  --max_timestep_ratio 0.96
+  --dfake_gen_update_ratio 5
+  --real_score_guidance_scale 3.0
+  # --real_score_guidance_scale_2 3.0
+  --fake_score_betas '0.0,0.999'
+  --warp_denoising_step
+)
+
+# Self-forcing specific arguments
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
+  --last_step_only False  # Whether to only use the last denoising step
+  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+  --validate_cache_structure False  # Set to True for debugging KV cache issues
+)
+
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+--nproc_per_node $NUM_GPUS \
+    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"
@@ -0,0 +1,180 @@
+#!/bin/bash
+#SBATCH --job-name=sf_distill_profile
+#SBATCH --partition=main
+#SBATCH --nodes=2
+#SBATCH --ntasks=2
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=sf_distill_output/sf_distill_profile.out
+#SBATCH --error=sf_distill_output/sf_distill_profile.err
+#SBATCH --exclusive
+set -e -x
+
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate wei-fv
+export HOME="/mnt/weka/home/hao.zhang/wei"
+
+# Basic Info
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29503
+export NODE_RANK=$SLURM_PROCID
+nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
+export MASTER_ADDR=${nodes[0]}
+export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
+export TOKENIZERS_PARALLELISM=false
+export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=offline
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+export FASTVIDEO_TORCH_PROFILER_DIR=/mnt/weka/home/hao.zhang/wei/FastVideo/traces
+export FASTVIDEO_TORCH_PROFILE_REGIONS=profiler_region_model_loading
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+# Configs
+NUM_GPUS=8
+NUM_TOTAL_GPUS=16
+
+# Model paths for Self-Forcing DMD distillation:
+GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"  # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Critic model
+
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
+# DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+# VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_2.json"
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  # Updated for self-forcing DMD
+  --output_dir "/mnt/sharefs/users/hao.zhang/wei/SFwan_t2v_finetune"
+  --wandb_run_name "sf_distill_2e-6_4e-7_4n"
+  # --use_sf_wan
+  # --sf_ode_init_path "checkpoints/ode_init.pt"
+  --max_train_steps 2
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 81  # Must be divisible by num_frame_per_block (81 % 3 = 0 ✓)
+  --enable_gradient_checkpointing_type "full"
+  --log_visualization
+  --simulate_generator_forward
+  --num_frame_per_block 3  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 21
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_TOTAL_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
+)
+
+# Model arguments
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH  # TODO: check if you can remove this in this script
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --generator_model_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 8
+)
+
+# Validation arguments
+validation_args=(
+#   --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 100
+  --validation_sampling_steps "4"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-6
+  --fake_score_learning_rate 4e-7
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 10.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 5
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 200
+  --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
+  # --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/FastVideo2/vidprom_8b16k_1e-5_gn1/checkpoint-3500/transformer/diffusion_pytorch_model.safetensors"
+#   --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/sf2/checkpoint_v3_7k/model.safetensors"
+)
+
+# Self-forcing DMD arguments
+dmd_args=(
+  --dmd_denoising_steps '1000,750,500,250'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --dfake_gen_update_ratio 5
+  --real_score_guidance_scale 3.0
+  --fake_score_betas '0.0,0.999'
+  --warp_denoising_step
+)
+
+# Self-forcing specific arguments
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
+  --last_step_only False  # Whether to only use the last denoising step
+  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+  --validate_cache_structure False  # Set to True for debugging KV cache issues
+)
+
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+--nproc_per_node $NUM_GPUS \
+    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"
@@ -0,0 +1,183 @@
+#!/bin/bash
+#SBATCH --job-name=moe_8n_sf_distill
+#SBATCH --partition=main
+#SBATCH --nodes=8
+#SBATCH --ntasks=8
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=moe_8n_sf_distill_output/moe_sf_distill_%j.out
+#SBATCH --error=moe_8n_sf_distill_output/moe_sf_distill_%j.err
+#SBATCH --exclusive
+
+# Basic Info
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export NCCL_DEBUG_SUBSYS=INIT,NET
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29500
+export NODE_RANK=$SLURM_PROCID
+nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
+export MASTER_ADDR=${nodes[0]}
+export TOKENIZERS_PARALLELISM=false
+# export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde"
+export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+# export WANDB_API_KEY='your_wandb_api_key_here'
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+source ~/conda/miniconda/bin/activate
+conda activate wei-fv
+export HOME="/mnt/weka/home/hao.zhang/wei"
+
+# Configs
+NUM_GPUS=8
+NUM_TOTAL_GPUS=64
+
+# Model paths for Self-Forcing DMD distillation with Wan2.2:
+# GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+GENERATOR_MODEL_PATH="rand0nmr/SFWan2.2-T2V-A14B-Diffusers"  # Updated to Wan2.2
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model
+# GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" # Updated to Wan2.2
+# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Teacher model
+# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model
+
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+# DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
+DATA_DIR="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/vidprom_16k_text_embed"
+# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset
+# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+training_args=(
+  --tracker_project_name SFwan2.2_t2v_distill_self_forcing_dmd  # Updated for Wan2.2
+  # --output_dir "/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_x0_cache_detach_df3000_cfg2_31_split"
+  --output_dir "/mnt/weka/home/hao.zhang/SFwan2.2_x0odeinit_x0CacheAppend_noaddnoisehigh_22cfg_moecritic_4steps"
+  # --output_dir "/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_x0_CacheAppend_df3000_cfg2"
+  --wandb_run_name "x0odeinit_x0CacheAppend_noaddnoisehigh_22cfg_moecritic_4steps"
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480  # Updated to match Wan2.2 config
+  --num_width 832   # Updated to match Wan2.2 config
+  --enable_gradient_checkpointing_type "full"
+  --simulate_generator_forward
+  --log_visualization
+  --num_frames 81
+  --num_frame_per_block 3  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 21
+  # --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_inference_transformer
+  # --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_2_inference_transformer
+  # --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wl/sharefs/wl/release/wan2.2-moe-ode-init/official_ode_init_wan22_high_ckpt2000_sharded_new
+  # --init_weights_from_safetensors_2 /mnt/weka/home/hao.zhang/wl/sharefs/wl/release/wan2.2-moe-ode-init/official_ode_init_wan22_low_ckpt2000_sharded_new
+  --init_weights_from_safetensors /mnt/weka/home/hao.zhang/official_ode_init_wan22_high_xbx0_4steps/checkpoint_model_002000/model.safetensors
+  --init_weights_from_safetensors_2 /mnt/weka/home/hao.zhang/official_ode_init_wan22_low_4steps/checkpoint_model_002000/model.safetensors
+  # --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-736_weight_only/generator_inference_transformer
+  # --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-736_weight_only/generator_2_inference_transformer
+  # --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors
+  # --init_weights_from_safetensors "/mnt/sharefs/users/hao.zhang/wei/diffusion_init_4000_14b_low_dense_shift5/model.safetensors"
+)
+
+parallel_args=(
+  --num_gpus $NUM_TOTAL_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
+)
+
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 10
+  --validation_sampling_steps "4"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+  --vae_cpu_offload True
+)
+
+optimizer_args=(
+  --learning_rate 2e-6
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 49
+  --weight_only_checkpointing_steps 99999
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 10.0
+)
+
+miscellaneous_args=(
+  # --use_add_noise_high
+  # --override_boundary_ratio 0.866
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 12
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 5
+  # --resume_from_checkpoint "/mnt/weka/home/hao.zhang/SFwan2.2_x0odeinit_x0CacheAppend_noaddnoisehigh_22cfg_moecritic/checkpoint-441"
+)
+
+dmd_args=(
+  # --dmd_denoising_steps '1000,850,800,750,570,450,350,250'
+  # --dmd_denoising_steps '1000,850,700,550,350,275,200,125'
+  --dmd_denoising_steps '1000,550,350,125'
+  # --dmd_denoising_steps '1000,750,500,250'
+  --min_timestep_ratio 0.04
+  --max_timestep_ratio 0.96
+  --dfake_gen_update_ratio 5
+  --real_score_guidance_scale 2.0
+  --real_score_guidance_scale_2 2.0
+  --fake_score_learning_rate 4e-7
+  --fake_score_betas '0.0,0.999'
+  --warp_denoising_step
+)
+
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
+  --last_step_only False  # Whether to only use the last denoising step
+  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+)
+
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
+--nproc_per_node $NUM_GPUS \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"
@@ -0,0 +1,178 @@
+#!/bin/bash
+#SBATCH --job-name=dmd_1333
+#SBATCH --partition=main
+#SBATCH --nodes=8
+#SBATCH --ntasks=8
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=dmd_1333_output/dmd_1333_%j.out
+#SBATCH --error=dmd_1333_output/dmd_1333_%j.err
+#SBATCH --exclusive
+
+# Basic Info
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export NCCL_DEBUG_SUBSYS=INIT,NET
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29500
+export NODE_RANK=$SLURM_PROCID
+nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
+export MASTER_ADDR=${nodes[0]}
+export TOKENIZERS_PARALLELISM=false
+# export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde"
+export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+# export WANDB_API_KEY='your_wandb_api_key_here'
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+source ~/conda/miniconda/bin/activate
+conda activate wei-fv
+export HOME="/mnt/weka/home/hao.zhang/wei"
+
+# Configs
+NUM_GPUS=8
+NUM_TOTAL_GPUS=64
+
+# Model paths for Self-Forcing DMD distillation with Wan2.2:
+# GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+GENERATOR_MODEL_PATH="rand0nmr/SFWan2.2-T2V-A14B-Diffusers"  # Updated to Wan2.2
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model
+# GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" # Updated to Wan2.2
+# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Teacher model
+# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model
+
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+# DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
+DATA_DIR="/mnt/sharefs/users/hao.zhang/wl/release/vidprom_16k_text_embed"
+# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset
+# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+training_args=(
+  --tracker_project_name SFwan2.2_t2v_distill_self_forcing_dmd  # Updated for Wan2.2
+  # --output_dir "/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_x0_cache_detach_df3000_cfg2_31_split"
+  --output_dir "/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_x0_CacheAppend_1333_3a_4a_no_add_noise_high"
+  # --output_dir "/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_x0_CacheAppend_df3000_cfg2"
+  --wandb_run_name "x0_CacheAppend_1333_3a_4a_no_add_noise_high"
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 19
+  --num_height 480  # Updated to match Wan2.2 config
+  --num_width 832   # Updated to match Wan2.2 config
+  --enable_gradient_checkpointing_type "full"
+  --simulate_generator_forward
+  --log_visualization
+  --num_frames 73
+  --num_frame_per_block 3  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 21
+  # --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_inference_transformer
+  # --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_2_inference_transformer
+  --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wei/diffusion_init_high_2800_shift5_1333/model.safetensors 
+  --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wei/diffusion_init_low_2800_shift5_1333/model.safetensors
+  # --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-736_weight_only/generator_inference_transformer
+  # --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-736_weight_only/generator_2_inference_transformer
+  # --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors
+  # --init_weights_from_safetensors "/mnt/sharefs/users/hao.zhang/wei/diffusion_init_4000_14b_low_dense_shift5/model.safetensors"
+)
+
+parallel_args=(
+  --num_gpus $NUM_TOTAL_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
+)
+
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "8"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+  # --vae_cpu_offload True
+)
+
+optimizer_args=(
+  --learning_rate 2e-6
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 49
+  --weight_only_checkpointing_steps 9999999
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 10.0
+)
+
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 12
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 197
+  --resume_from_checkpoint "/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_x0_CacheAppend_1333_3a_4a_no_add_noise_high/checkpoint-196"
+)
+
+dmd_args=(
+  --dmd_denoising_steps '1000,850,700,550,350,275,200,125'
+  # --dmd_denoising_steps '1000,850,350,275'
+  # --dmd_denoising_steps '1000,750,500,250'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --dfake_gen_update_ratio 5
+  --real_score_guidance_scale 2.0
+  --real_score_guidance_scale_2 2.0
+  --fake_score_learning_rate 4e-7
+  --fake_score_betas '0.0,0.999'
+  --warp_denoising_step
+)
+
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
+  --last_step_only False  # Whether to only use the last denoising step
+  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+)
+
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
+--nproc_per_node $NUM_GPUS \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"
@@ -0,0 +1,144 @@
+#!/bin/bash
+#SBATCH --job-name=moe_dmd_distill
+#SBATCH --partition=main
+#SBATCH --nodes=2
+#SBATCH --ntasks=2
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=moe_dmd_distill_output/moe_dmd_distill_%j.out
+#SBATCH --error=moe_dmd_distill_output/moe_dmd_distill_%j.err
+#SBATCH --exclusive
+
+# Basic Info
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export NCCL_DEBUG_SUBSYS=INIT,NET
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29500
+export NODE_RANK=$SLURM_PROCID
+nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
+export MASTER_ADDR=${nodes[0]}
+export TOKENIZERS_PARALLELISM=false
+export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+# export WANDB_API_KEY='your_wandb_api_key_here'
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+source ~/conda/miniconda/bin/activate
+conda activate wei-fv
+export HOME="/mnt/weka/home/hao.zhang/wei"
+
+# Configs
+NUM_GPUS=8
+NUM_TOTAL_GPUS=16
+
+# Model paths for Self-Forcing DMD distillation with Wan2.2:
+GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"  # Updated to Wan2.2
+# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model
+# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model
+# GENERATOR_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Updated to Wan2.2
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model
+
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+DATA_DIR="/mnt/weka/home/hao.zhang/wei/FastVideo/data/crush-smol_processed_t2v/combined_parquet_dataset"
+# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset
+# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+training_args=(
+  --tracker_project_name SFwan2.2_t2v_distill_dmd  # Updated for Wan2.2
+  --output_dir "/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_dmd"
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 448  # Updated to match Wan2.2 config
+  --num_width 832   # Updated to match Wan2.2 config
+  --enable_gradient_checkpointing_type "full"
+  # --log_visualization
+  --simulate_generator_forward
+  --num_frames 81
+  # --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/high/3k/
+  # --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/low/3k/
+)
+
+parallel_args=(
+  --num_gpus $NUM_TOTAL_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
+)
+
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 20
+  --validation_sampling_steps "4"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+optimizer_args=(
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --max_grad_norm 1.0
+)
+
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 5
+  --seed 1000
+  --ema_start_step 100
+)
+
+dmd_args=(
+  --dmd_denoising_steps '1000,750,500,250'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --generator_update_interval 5
+  --real_score_guidance_scale 3.0
+)
+
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
+--nproc_per_node $NUM_GPUS \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+    fastvideo/training/wan_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}"
@@ -0,0 +1,131 @@
+#!/bin/bash
+#SBATCH --job-name=moe_finetune
+#SBATCH --partition=main
+#SBATCH --nodes=2
+#SBATCH --ntasks=2
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=moe_output/moe_%j.out
+#SBATCH --error=moe_output/moe_%j.err
+#SBATCH --exclusive
+
+# Basic Info
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export NCCL_DEBUG_SUBSYS=INIT,NET
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29500
+export NODE_RANK=$SLURM_PROCID
+nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
+export MASTER_ADDR=${nodes[0]}
+export TOKENIZERS_PARALLELISM=false
+export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+# export WANDB_API_KEY='your_wandb_api_key_here'
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+source ~/conda/miniconda/bin/activate
+conda activate wei-fv
+export HOME="/mnt/weka/home/hao.zhang/wei"
+
+
+# Configs
+NUM_GPUS=8
+NUM_TOTAL_GPUS=16
+
+# Model paths for Self-Forcing DMD distillation with Wan2.2:
+GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"  # Updated to Wan2.2
+# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model
+# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model
+
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+DATA_DIR="/mnt/weka/home/hao.zhang/wei/FastVideo/data/crush-smol_processed_t2v/combined_parquet_dataset"
+# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset
+# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wei/FastVideo/data/crush-smol-single_processed_t2v/validation.json"
+
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+training_args=(
+  --tracker_project_name VSA_finetune  # Updated for Wan2.2
+  --output_dir "/mnt/sharefs/users/hao.zhang/wei/Wan2.2-MoE-finetune"
+  --max_train_steps 200
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 448  # Updated to match Wan2.2 config
+  --num_width 832   # Updated to match Wan2.2 config
+  --enable_gradient_checkpointing_type "full"
+  # --log_visualization
+  --num_frames 81
+  # --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/high/3k/
+  # --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/low/3k/
+)
+
+parallel_args=(
+  --num_gpus $NUM_TOTAL_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
+)
+
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+)
+
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+validation_args=(
+#   --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "40"
+  --validation_guidance_scale "5.0" # not used for dmd inference
+)
+
+optimizer_args=(
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --dit_cpu_offload True
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 1.0
+)
+
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 5
+  --seed 1000
+)
+
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
+--nproc_per_node $NUM_GPUS \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+    fastvideo/training/wan_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"
@@ -7,6 +7,8 @@ from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
 def is_blocks(n: str, m) -> bool:
    return "blocks" in n and str.isdigit(n.split(".")[-1])

+# def is_not_blocks(n: str, m) -> bool:
+#     return "blocks" not in n and len(n.split(".")) == 1

@dataclass
 class WanVideoArchConfig(DiTArchConfig):
@@ -99,7 +101,9 @@ class WanVideoArchConfig(DiTArchConfig):
    local_attn_size: int = -1  # Window size for temporal local attention (-1 indicates global attention)
    sink_size: int = 0  # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache
    num_frames_per_block: int = 3
+    # sliding_window_num_frames: int = 19
    sliding_window_num_frames: int = 21
+    # sliding_window_num_frames: int = 40

    def __post_init__(self):
        super().__post_init__()
@@ -92,6 +92,9 @@ class PipelineConfig:
    # Compilation
    # enable_torch_compile: bool = False

+    # Self-Forcing parameters
+    is_causal: bool = False
+
    @staticmethod
    def add_cli_args(parser: FlexibleArgumentParser,
                     prefix: str = "") -> FlexibleArgumentParser:
@@ -13,7 +13,7 @@ from fastvideo.configs.pipelines.wan import (
    FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config,
    Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config,
    WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig,
-    SelfForcingWanT2V480PConfig)
+    SelfForcingWanT2V480PConfig, SelfForcingMoEWanT2V480PConfig, SelfForcingMoEWanT2V720PConfig)
 # isort: on
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
@@ -36,6 +36,9 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
    "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig,
    "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
+    "rand0nmr/SFWan2.1-T2V-A1.3B-Diffusers": SelfForcingMoEWanT2V480PConfig,
+    "rand0nmr/SFWan2.2-T2V-A14B-Diffusers": SelfForcingMoEWanT2V720PConfig,
+    "rand0nmr/SFWan2.2-T2V-14B-Diffusers": SelfForcingWanT2V480PConfig,
    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_Config,
    "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_Config,
    "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_Config,
@@ -164,4 +164,36 @@ class SelfForcingWanT2V480PConfig(WanT2V480PConfig):
    flow_shift: float | None = 5.0
    dmd_denoising_steps: list[int] | None = field(
        default_factory=lambda: [1000, 750, 500, 250])
+    # flow_shift: float | None = 12.0
+    # dmd_denoising_steps: list[int] | None = field(
+    #     default_factory=lambda: [1000,350,275,200,125])
    warp_denoising_step: bool = True
+
+    def __post_init__(self) -> None:
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+
+@dataclass
+class SelfForcingMoEWanT2V480PConfig(SelfForcingWanT2V480PConfig):
+    boundary_ratio: float | None = 0.875
+
+    def __post_init__(self) -> None:
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+        self.dit_config.boundary_ratio = self.boundary_ratio
+
+@dataclass
+class SelfForcingMoEWanT2V720PConfig(SelfForcingMoEWanT2V480PConfig):
+    flow_shift: float | None = 12.0
+    # dmd_denoising_steps: list[int] | None = field(
+    #     default_factory=lambda: [1000,850,700,550,350,275,200,125])
+    dmd_denoising_steps: list[int] | None = field(
+        default_factory=lambda: [1000,550,350,125])
+    # flow_shift: float | None = 5.0
+    # dmd_denoising_steps: list[int] | None = field(
+    #     default_factory=lambda: [1000,850,800,750,570,450,350,250])
+
+    def __post_init__(self) -> None:
+        self.dit_config.boundary_ratio = self.boundary_ratio
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
@@ -19,6 +19,7 @@ from fastvideo.configs.sample.wan import (
    WanT2V_1_3B_SamplingParam,
    WanT2V_14B_SamplingParam,
    SelfForcingWanT2V480PConfig,
+    SelfForcingMoEWanT2V480PConfig,
 )
 # isort: on
 from fastvideo.logger import init_logger
@@ -55,6 +56,11 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {

    # Causal Self-Forcing Wan2.1
    "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
+    "rand0nmr/SFWan2.2-T2V-14B-Diffusers": SelfForcingWanT2V480PConfig,
+    "rand0nmr/SFWan2.1-T2V-A1.3B-Diffusers": SelfForcingWanT2V480PConfig,
+
+    # Causal Self-Forcing Wan2.2 MoE
+    "rand0nmr/SFWan2.2-T2V-A14B-Diffusers": SelfForcingMoEWanT2V480PConfig,
    # Add other specific weight variants
 }

@@ -167,4 +167,17 @@ class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam):
 # =============================================
@dataclass
 class SelfForcingWanT2V480PConfig(WanT2V_1_3B_SamplingParam):
-    pass
+    height: int = 480
+    width: int = 832
+    num_frames: int = 81
+    fps: int = 16
+
+@dataclass
+class SelfForcingMoEWanT2V480PConfig(WanT2V_1_3B_SamplingParam):
+    height: int = 480
+    width: int = 832
+    # height: int = 720
+    # width: int = 1280
+    # num_frames: int = 157
+    num_frames: int = 81
+    fps: int = 16
@@ -166,6 +166,8 @@ class FastVideoArgs:

    # MoE parameters used by Wan2.2
    boundary_ratio: float | None = 0.875
+    override_boundary_ratio: float | None = None
+    use_add_noise_high: bool = False

    @property
    def training_mode(self) -> bool:
@@ -411,6 +413,18 @@ class FastVideoArgs:
            default=FastVideoArgs.override_transformer_cls_name,
            help="Override transformer cls name",
        )
+        parser.add_argument(
+            "--override-boundary-ratio",
+            type=float,
+            default=FastVideoArgs.override_boundary_ratio,
+            help="Override boundary ratio for MoE",
+        )
+        parser.add_argument(
+            "--use-add-noise-high",
+            action=StoreBoolean,
+            default=FastVideoArgs.use_add_noise_high,
+            help="Whether to use add_noise_high & pred_noise_to_x_bound for MoE",
+        )
        # Add pipeline configuration arguments
        PipelineConfig.add_cli_args(parser)

@@ -717,6 +731,7 @@ class TrainingArgs(FastVideoArgs):
    min_timestep_ratio: float = 0.2
    max_timestep_ratio: float = 0.98
    real_score_guidance_scale: float = 3.5
+    real_score_guidance_scale_2: float = 3.5
    fake_score_learning_rate: float = 0.0  # separate learning rate for fake_score_transformer, if 0.0, use learning_rate
    fake_score_lr_scheduler: str = "constant"  # separate lr scheduler for fake_score_transformer, if not set, use lr_scheduler
    fake_score_betas: str = "0.9,0.999"  # betas for fake score optimizer, format: "beta1,beta2"
@@ -1102,6 +1117,10 @@ class TrainingArgs(FastVideoArgs):
                            type=float,
                            default=TrainingArgs.real_score_guidance_scale,
                            help="Teacher guidance scale")
+        parser.add_argument("--real-score-guidance-scale-2",
+                            type=float,
+                            default=TrainingArgs.real_score_guidance_scale_2,
+                            help="Teacher guidance scale")
        parser.add_argument("--fake-score-learning-rate",
                            type=float,
                            default=TrainingArgs.fake_score_learning_rate,
@@ -2,10 +2,10 @@

 import math
 from typing import Any
-
+import gc
 import numpy as np
-import torch
-import torch.nn as nn
+import torch  # ty:ignore[unresolved-import]
+import torch.nn as nn  # pyright: ignore[reportMissingImports]  # ty:ignore[unresolved-import]

 from torch.nn.attention.flex_attention import create_block_mask, flex_attention
 from torch.nn.attention.flex_attention import BlockMask
@@ -20,7 +20,7 @@ import fastvideo.envs as envs
 from fastvideo.attention import (DistributedAttention,
                                 LocalAttention)
 from fastvideo.configs.models.dits import WanVideoConfig
-from fastvideo.distributed.parallel_state import get_sp_world_size
+from fastvideo.distributed.parallel_state import get_sp_world_size, get_local_torch_device
 from fastvideo.forward_context import get_forward_context
 from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift,
                                        RMSNorm, ScaleResidual,
@@ -36,6 +36,26 @@ from fastvideo.models.dits.wanvideo import WanT2VCrossAttention, WanTimeTextImag
 from fastvideo.platforms import AttentionBackendEnum, current_platform

 logger = init_logger(__name__)
+
+class CacheAppend(torch.autograd.Function):
+    """
+    KV cache with shape [batch, seq_len, heads, head_dim].
+    """
+    @staticmethod
+    def forward(ctx, storage, active_cache, x, start, end):
+        # Ensure storage has the same dtype as x
+        storage.data[:, start:end] = x
+        ctx.save_for_backward(storage.to(x.dtype))
+        ctx.start = start
+        ctx.end = end
+        return storage[:, :end].to(x.dtype) # Ensure returned value has same dtype as input
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        start = ctx.start
+        end = ctx.end
+        return None, grad_output[:, :start], grad_output[:, start:end], None, None
+
 class CausalWanSelfAttention(nn.Module):

    def __init__(self,
@@ -57,6 +77,9 @@ class CausalWanSelfAttention(nn.Module):
        self.eps = eps
        self.parallel_attention = parallel_attention
        self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 1560
+        # self.max_attention_size = 29640 # 1560 * 19
+        # self.max_attention_size = 3600 * 19 # 720P
+        # self.max_attention_size = 40 * 1560 # 10 seconds

        # Scaled dot product attention
        self.attn = LocalAttention(
@@ -67,6 +90,9 @@ class CausalWanSelfAttention(nn.Module):
            causal=False,
            supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
                                          AttentionBackendEnum.TORCH_SDPA))
+        
+        self.k_cache = None
+        self.v_cache = None

    def forward(self, 
                q: torch.Tensor,
@@ -84,6 +110,33 @@ class CausalWanSelfAttention(nn.Module):
            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
        """
+        # def drop_cache(grad):
+        #     # self.k_cache = torch.cat([self.k_cache[:, :current_start], self.k_cache[:, current_start:].detach()], dim=1)
+        #     # self.v_cache = torch.cat([self.v_cache[:, :current_start], self.v_cache[:, current_start:].detach()], dim=1)
+        #     self.k_cache = self.k_cache[:, :current_start]
+        #     self.v_cache = self.v_cache[:, :current_start]
+        #     if self.k_cache.numel() == 0:
+        #         self.k_cache = None
+        #     if self.v_cache.numel() == 0:
+        #         self.v_cache = None
+        #     gc.collect()
+        #     torch.cuda.empty_cache()
+        #     return grad
+
+        # if q.requires_grad:
+        #     q.register_hook(drop_cache)
+
+        if kv_cache is not None:
+            # Then we are in _forward_inference mode
+            if self.k_cache is None:
+                del self.k_cache
+                self.register_buffer("k_cache", torch.empty(1, self.max_attention_size, self.num_heads, self.head_dim, device=q.device, dtype=q.dtype), persistent=False)
+                self.k_cache.requires_grad_(True)
+            if self.v_cache is None:
+                del self.v_cache
+                self.register_buffer("v_cache", torch.empty(1, self.max_attention_size, self.num_heads, self.head_dim, device=v.device, dtype=v.dtype), persistent=False)
+                self.v_cache.requires_grad_(True)
+
        if cache_start is None:
            cache_start = current_start

@@ -128,35 +181,55 @@ class CausalWanSelfAttention(nn.Module):
            num_new_tokens = roped_query.shape[1]
            if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
                    num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
+                # @TODO(Wei): This part has not been thoroughly tested yet. Use with caution.
                # Calculate the number of new tokens added in this step
                # Shift existing cache content left to discard oldest tokens
                # Clone the source slice to avoid overlapping memory error
                num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
                num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
-                kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                    kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-                kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                    kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+                # kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+                #     kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+                # kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+                #     kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+                self.k_cache[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+                    self.k_cache[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+                self.v_cache[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+                    self.v_cache[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
                # Insert the new keys/values at the end
                local_end_index = kv_cache["local_end_index"].item() + current_end - \
                    kv_cache["global_end_index"].item() - num_evicted_tokens
                local_start_index = local_end_index - num_new_tokens
-                kv_cache["k"][:, local_start_index:local_end_index] = roped_key
-                kv_cache["v"][:, local_start_index:local_end_index] = v
+                local_k = CacheAppend.apply(self.k_cache, kv_cache["k"], roped_key, local_start_index, local_end_index)
+                local_v = CacheAppend.apply(self.v_cache, kv_cache["v"], v, local_start_index, local_end_index)
+                # kv_cache["k"][:, local_start_index:local_end_index] = roped_key
+                # kv_cache["v"][:, local_start_index:local_end_index] = v
            else:
+                assert kv_cache["local_end_index"].item() == kv_cache["global_end_index"].item(), "local_end_index and global_end_index must be the same"
                # Assign new keys/values directly up to current_end
                local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
                local_start_index = local_end_index - num_new_tokens
-                kv_cache["k"] = kv_cache["k"].detach()
-                kv_cache["v"] = kv_cache["v"].detach()
-                # logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None)
-                kv_cache["k"][:, local_start_index:local_end_index] = roped_key
-                kv_cache["v"][:, local_start_index:local_end_index] = v
+                # kv_cache["k"] = kv_cache["k"].detach()
+                # kv_cache["v"] = kv_cache["v"].detach()
+                # kv_cache["k"][:, local_start_index:local_end_index] = roped_key
+                # kv_cache["v"][:, local_start_index:local_end_index] = v
+                local_k = CacheAppend.apply(self.k_cache, kv_cache["k"], roped_key, local_start_index, local_end_index)
+                # logger.info("Is local_k meta tensor: %s, %s", local_k.is_meta, local_k.shape)
+                # logger.info("local_start_index: %d, local_end_index: %d, number of zeros in local k: %d", local_start_index, local_end_index, (local_k == 0).sum().item())
+                local_v = CacheAppend.apply(self.v_cache, kv_cache["v"], v, local_start_index, local_end_index)
+
            x = self.attn(
                roped_query,
-                kv_cache["k"][:, max(0, local_end_index - self.max_attention_size):local_end_index],
-                kv_cache["v"][:, max(0, local_end_index - self.max_attention_size):local_end_index]
+                local_k,
+                local_v
            )
+            # x = self.attn(
+            #     roped_query,
+            #     kv_cache["k"][:, max(0, local_end_index - self.max_attention_size):local_end_index],
+            #     kv_cache["v"][:, max(0, local_end_index - self.max_attention_size):local_end_index]
+            # )
+
+            kv_cache["k"] = local_k
+            kv_cache["v"] = local_v
            kv_cache["global_end_index"].fill_(current_end)
            kv_cache["local_end_index"].fill_(local_end_index)

@@ -233,6 +306,9 @@ class CausalWanTransformerBlock(nn.Module):

        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)

+        self.null_shift = torch.tensor([0], device=get_local_torch_device())
+        self.null_scale = torch.tensor([0], device=get_local_torch_device())
+
    def forward(
        self,
        hidden_states: torch.Tensor,
@@ -283,9 +359,8 @@ class CausalWanTransformerBlock(nn.Module):
        attn_output, _ = self.to_out(attn_output)
        attn_output = attn_output.squeeze(1)

-        null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
        norm_hidden_states, hidden_states = self.self_attn_residual_norm(
-            hidden_states, attn_output, gate_msa, null_shift, null_scale)
+            hidden_states, attn_output, gate_msa, self.null_shift, self.null_scale)

        # 2. Cross-attention
        attn_output = self.attn2(norm_hidden_states,
@@ -293,6 +293,8 @@ class TextEncoderLoader(ComponentLoader):
                        mesh=mesh["offload"],
                        fsdp_shard_conditions=model._fsdp_shard_conditions,
                        pin_cpu_memory=fastvideo_args.pin_cpu_memory)
+            else:
+                raise ValueError("text_encoder_cpu_offload must be True")
            # We only enable strict check for non-quantized models
            # that have loaded weights tracking currently.
            # if loaded_weights is not None:
@@ -507,8 +509,9 @@ class SchedulerLoader(ComponentLoader):
        scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name)

        scheduler = scheduler_cls(**config)
-        if fastvideo_args.pipeline_config.flow_shift is not None:
-            scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift)
+        assert config['shift'] == 12, "shift must be 12"
+        # if fastvideo_args.pipeline_config.flow_shift is not None:
+        #     scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift)
        if fastvideo_args.pipeline_config.timesteps_scale is not None:
            scheduler.set_timesteps_scale(
                fastvideo_args.pipeline_config.timesteps_scale)
@@ -91,6 +91,7 @@ def maybe_load_fsdp_model(
    )

    logger.info("Loading model with default_dtype: %s", default_dtype)
+    logger.info("Loading model with cpu_offload: %s", cpu_offload)
    with set_default_dtype(default_dtype), torch.device("meta"):
        model = model_cls(**init_params)

@@ -86,6 +86,12 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
            return (prev_sample, )
        return SelfForcingFlowMatchSchedulerOutput(prev_sample=prev_sample)

+    @staticmethod
+    def calculate_alpha_beta_high(sigma, sigma_bound):
+        alpha = (1 - sigma) / (1 - sigma_bound)
+        beta = torch.sqrt(sigma ** 2 - (alpha * sigma_bound) ** 2)
+        return alpha, beta
+
    def add_noise(self, original_samples, noise, timestep):
        """
        Diffusion forward corruption process.
@@ -105,6 +111,32 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
        sample = (1 - sigma) * original_samples + sigma * noise
        return sample.type_as(noise)

+    def add_noise_high(self, original_samples, noise, timestep, boundary_timestep):
+        """
+        Diffusion forward corruption process.
+        Input:
+            - clean_latent: the clean latent with shape [B*T, C, H, W]
+            - noise: the noise with shape [B*T, C, H, W]
+            - timestep: the timestep with shape [B*T]
+        Output: the corrupted latent with shape [B*T, C, H, W]
+        """
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        if boundary_timestep.ndim == 2:
+            boundary_timestep = boundary_timestep.flatten(0, 1)
+        self.sigmas = self.sigmas.to(noise.device)
+        self.timesteps = self.timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+
+        boundary_timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1)
+        sigma_boundary = self.sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1)
+        alpha, beta = self.calculate_alpha_beta_high(sigma, sigma_boundary)
+        sample = alpha * original_samples + beta * noise
+        return sample.type_as(noise)
+
    def training_target(self, sample, noise, timestep):
        target = noise - sample
        return target
@@ -180,3 +180,51 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor,
    sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
    pred_video = noise_input_latent - sigma_t * pred_noise
    return pred_video.to(dtype)
+
+def pred_noise_to_x_bound(pred_noise: torch.Tensor,
+                             noise_input_latent: torch.Tensor,
+                             timestep: torch.Tensor,
+                             boundary_timestep: torch.Tensor,
+                             scheduler: Any) -> torch.Tensor:
+    """
+    Convert predicted noise to clean latent.
+
+    Args:
+    pred_noise: the predicted noise with shape [B, C, H, W]
+        where B is batch_size or batch_size * num_frames
+    noise_input_latent: the noisy latent with shape [B, C, H, W],
+    timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames]
+    boundary_timestep: the boundary timestep with shape [1] or [bs * num_frames] or [bs, num_frames]
+    scheduler: the scheduler
+
+    Returns:
+        the predicted video with shape [B, C, H, W]
+    """
+    # If timestep is [bs, num_frames]
+    if timestep.ndim == 2:
+        timestep = timestep.flatten(0, 1)
+        assert timestep.numel() == noise_input_latent.shape[0]
+    elif timestep.ndim == 1:
+        # If timestep is [1]
+        if timestep.shape[0] == 1:
+            timestep = timestep.expand(noise_input_latent.shape[0])
+        else:
+            assert timestep.numel() == noise_input_latent.shape[0]
+    else:
+        raise ValueError(f"[pred_noise_to_pred_video] Invalid timestep shape: {timestep.shape}")
+    # timestep shape should be [B]
+    dtype = pred_noise.dtype
+    device = pred_noise.device
+    pred_noise = pred_noise.double().to(device)
+    noise_input_latent = noise_input_latent.double().to(device)
+    sigmas = scheduler.sigmas.double().to(device)
+    timesteps = scheduler.timesteps.double().to(device)
+    timestep_id = torch.argmin(
+        (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+    sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+
+    boundary_timestep_id = torch.argmin(
+        (timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1)
+    sigma_t_boundary = sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1)
+    pred_video = noise_input_latent - (sigma_t - sigma_t_boundary) * pred_noise
+    return pred_video.to(dtype)
@@ -50,7 +50,8 @@ class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
                       stage=CausalDMDDenosingStage(
                           transformer=self.get_module("transformer"),
                           transformer_2=self.get_module("transformer_2", None),
-                           scheduler=self.get_module("scheduler")))
+                           scheduler=self.get_module("scheduler"),
+                           vae=self.get_module("vae")))

        self.add_stage(stage_name="decoding_stage",
                       stage=DecodingStage(vae=self.get_module("vae")))
@@ -71,6 +71,7 @@ class ComposedPipelineBase(ABC):

        # Load modules directly in initialization
        logger.info("Loading pipeline modules...")
+        fastvideo_args.dit_cpu_offload = True
        self.modules = self.load_modules(fastvideo_args, loaded_modules)

    def set_trainable(self) -> None:
@@ -234,6 +234,8 @@ class TrainingBatch:
    # Training outputs
    total_loss: float | None = None
    grad_norm: float | None = None
+    generator_grad_norm: float | None = None
+    critic_grad_norm: float | None = None

    # Distillation-specific attributes
    encoder_hidden_states_neg: torch.Tensor | None = None
@@ -1,10 +1,11 @@
 import torch  # type: ignore
+import gc

 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
-from fastvideo.models.utils import pred_noise_to_pred_video
+from fastvideo.models.utils import pred_noise_to_pred_video, pred_noise_to_x_bound
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.denoising import DenoisingStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
@@ -34,13 +35,14 @@ class CausalDMDDenosingStage(DenoisingStage):
    Denoising stage for causal diffusion.
    """

-    def __init__(self, transformer, scheduler, transformer_2=None) -> None:
-        super().__init__(transformer, scheduler, transformer_2)
+    def __init__(self, transformer, scheduler, transformer_2=None, vae=None) -> None:
+        super().__init__(transformer, scheduler, transformer_2, vae)
        # KV and cross-attention cache state (initialized on first forward)
        self.transformer = transformer
        self.transformer_2 = transformer_2
-        self.kv_cache1: list | None = None
-        self.crossattn_cache: list | None = None
+        self.vae = vae
+        # self.kv_cache1: list | None = None
+        # self.crossattn_cache: list | None = None
        # Model-dependent constants (aligned with causal_inference.py assumptions)
        self.num_transformer_blocks = len(self.transformer.blocks)
        self.num_frames_per_block = self.transformer.config.arch_config.num_frames_per_block
@@ -79,6 +81,7 @@ class CausalDMDDenosingStage(DenoisingStage):
                                                          dtype=torch.float32)))
            timesteps = scheduler_timesteps[1000 - timesteps]
        timesteps = timesteps.to(get_local_torch_device())
+        logger.info("timesteps: %s", timesteps)

        # Image kwargs (kept empty unless caller provides compatible args)
        image_kwargs: dict = {}
@@ -103,34 +106,43 @@ class CausalDMDDenosingStage(DenoisingStage):
        assert torch.isnan(prompt_embeds[0]).sum() == 0

        # Initialize or reset caches
-        if self.kv_cache1 is None:
-            self._initialize_kv_cache(batch_size=latents.shape[0],
+        kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0],
                                      dtype=target_dtype,
                                      device=latents.device)
-            self._initialize_crossattn_cache(
+        crossattn_cache = self._initialize_crossattn_cache(
                batch_size=latents.shape[0],
                max_text_len=fastvideo_args.pipeline_config.
                text_encoder_configs[0].arch_config.text_len,
                dtype=target_dtype,
                device=latents.device)
-        else:
-            assert self.crossattn_cache is not None
-            # reset cross-attention cache
-            for block_index in range(self.num_transformer_blocks):
-                self.crossattn_cache[block_index][
-                    "is_init"] = False  # type: ignore
-            # reset kv cache pointers
-            for block_index in range(len(self.kv_cache1)):
-                self.kv_cache1[block_index][
-                    "global_end_index"] = torch.tensor(  # type: ignore
-                        [0],
-                        dtype=torch.long,
-                        device=latents.device)
-                self.kv_cache1[block_index][
-                    "local_end_index"] = torch.tensor(  # type: ignore
-                        [0],
-                        dtype=torch.long,
-                        device=latents.device)
+        # if self.kv_cache1 is None:
+        #     self._initialize_kv_cache(batch_size=latents.shape[0],
+        #                               dtype=target_dtype,
+        #                               device=latents.device)
+        #     self._initialize_crossattn_cache(
+        #         batch_size=latents.shape[0],
+        #         max_text_len=fastvideo_args.pipeline_config.
+        #         text_encoder_configs[0].arch_config.text_len,
+        #         dtype=target_dtype,
+        #         device=latents.device)
+        # else:
+        #     assert self.crossattn_cache is not None
+        #     # reset cross-attention cache
+        #     for block_index in range(self.num_transformer_blocks):
+        #         self.crossattn_cache[block_index][
+        #             "is_init"] = False  # type: ignore
+        #     # reset kv cache pointers
+        #     for block_index in range(len(self.kv_cache1)):
+        #         self.kv_cache1[block_index][
+        #             "global_end_index"] = torch.tensor(  # type: ignore
+        #                 [0],
+        #                 dtype=torch.long,
+        #                 device=latents.device)
+        #         self.kv_cache1[block_index][
+        #             "local_end_index"] = torch.tensor(  # type: ignore
+        #                 [0],
+        #                 dtype=torch.long,
+        #                 device=latents.device)

        # Optional: cache context features from provided image latents prior to generation
        current_start_frame = 0
@@ -153,8 +165,8 @@ class CausalDMDDenosingStage(DenoisingStage):
                        image_first_btchw,
                        prompt_embeds,
                        t_zero,
-                        kv_cache=self.kv_cache1,
-                        crossattn_cache=self.crossattn_cache,
+                        kv_cache=kv_cache1,
+                        crossattn_cache=crossattn_cache,
                        current_start=current_start_frame *
                        self.frame_seq_length,
                        **image_kwargs,
@@ -179,8 +191,8 @@ class CausalDMDDenosingStage(DenoisingStage):
                        ref_btchw,
                        prompt_embeds,
                        t_zero,
-                        kv_cache=self.kv_cache1,
-                        crossattn_cache=self.crossattn_cache,
+                        kv_cache=kv_cache1,
+                        crossattn_cache=crossattn_cache,
                        current_start=current_start_frame *
                        self.frame_seq_length,
                        **image_kwargs,
@@ -193,36 +205,78 @@ class CausalDMDDenosingStage(DenoisingStage):
        pos_start_base = current_start_frame

        # Determine block sizes
-        if not independent_first_frame or (independent_first_frame
-                                           and batch.image_latent is not None):
-            if t % self.num_frames_per_block != 0:
-                raise ValueError(
-                    "num_frames must be divisible by num_frames_per_block for causal DMD denoising"
-                )
-            num_blocks = t // self.num_frames_per_block
-            block_sizes = [self.num_frames_per_block] * num_blocks
-            start_index = 0
+        block_sizes = [self.num_frames_per_block] * 7
+        start_index = 0
+
+        if fastvideo_args.override_boundary_ratio is not None:
+            boundary_timestep = fastvideo_args.override_boundary_ratio * self.scheduler.num_train_timesteps
+            logger.info("Overriding boundary timestep to %s", boundary_timestep)
        else:
-            if (t - 1) % self.num_frames_per_block != 0:
-                raise ValueError(
-                    "(num_frames - 1) must be divisible by num_frame_per_block when independent_first_frame=True"
-                )
-            num_blocks = (t - 1) // self.num_frames_per_block
-            block_sizes = [1] + [self.num_frames_per_block] * num_blocks
-            start_index = 0
+            if fastvideo_args.pipeline_config.dit_config.boundary_ratio is not None:
+                boundary_timestep = fastvideo_args.pipeline_config.dit_config.boundary_ratio * self.scheduler.num_train_timesteps
+            else:
+                boundary_timestep = None
+
+        high_noise_timesteps = timesteps[timesteps >= boundary_timestep]
+        assert len(high_noise_timesteps) == len(timesteps) // 2, f"There are incorrect number of {len(high_noise_timesteps)} high noise timesteps"
+
+        first_frame_latent = None
+        if batch.pil_image is not None:
+            # Causal video gen directly replaces the first frame of the latent with
+            # the image latent instead of appending along the channel dim
+            assert self.vae is not None, "VAE is not provided for causal video gen task"
+            self.vae = self.vae.to(get_local_torch_device())
+            first_frame_latent = self.vae.encode(batch.pil_image).mean.float()
+            if (hasattr(self.vae, "shift_factor")
+                    and self.vae.shift_factor is not None):
+                if isinstance(self.vae.shift_factor, torch.Tensor):
+                    first_frame_latent -= self.vae.shift_factor.to(first_frame_latent.device, first_frame_latent.dtype)
+                else:
+                    first_frame_latent -= self.vae.shift_factor
+
+            if isinstance(self.vae.scaling_factor, torch.Tensor):
+                first_frame_latent = first_frame_latent * self.vae.scaling_factor.to(first_frame_latent.device, first_frame_latent.dtype)
+            else:
+                first_frame_latent = first_frame_latent * self.vae.scaling_factor
+
+            # latents[:, :, :3, :, :] = first_frame_latent.unsqueeze(2).repeat(1, 1, 3, 1, 1)
+            # start_index += 3
+            # block_sizes.pop(0)
+            
+            if fastvideo_args.vae_cpu_offload:
+                self.vae = self.vae.to("cpu")
+            else:
+                raise ValueError("vae_cpu_offload must be True")
+
+            # if boundary_timestep is not None:
+                # first_frame_latent_bound = self.scheduler.add_noise(
+                #     first_frame_latent,
+                #     torch.randn_like(first_frame_latent),
+                #     boundary_timestep * torch.ones([1], device=first_frame_latent.device, dtype=torch.long))

        # DMD loop in causal blocks
        with self.progress_bar(total=len(block_sizes) *
                               len(timesteps)) as progress_bar:
-            for current_num_frames in block_sizes:
+            for block_idx, current_num_frames in enumerate(block_sizes):
                current_latents = latents[:, :, start_index:start_index +
                                          current_num_frames, :, :]
+
+                if block_idx == 0 and first_frame_latent is not None:
+                    # previous_frame = latents[:, :, start_index - 1:start_index, :, :].clone()
+                    first_frame_mask = torch.ones_like(current_latents)
+                    first_frame_mask[:, :, 0] = 0
+                    current_latents = (1 - first_frame_mask) * first_frame_latent + first_frame_mask * current_latents
+                    # if boundary_timestep is not None:
+                    #     current_latents = first_frame_latent_bound.repeat(1, 1, current_num_frames, 1, 1)
+                    # else:
+                    #     current_latents = first_frame_latent.repeat(1, 1, current_num_frames, 1, 1)
+                
                # use BTCHW for DMD conversion routines
                noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4)
                video_raw_latent_shape = noise_latents_btchw.shape

                for i, t_cur in enumerate(timesteps):
-                    if self.transformer_2 is not None and fastvideo_args.pipeline_config.dit_config.boundary_ratio is not None and t_cur < fastvideo_args.pipeline_config.dit_config.boundary_ratio * self.scheduler.num_train_timesteps:
+                    if boundary_timestep is not None and t_cur < boundary_timestep:
                        current_model = self.transformer_2
                    else:
                        current_model = self.transformer
@@ -273,15 +327,19 @@ class CausalDMDDenosingStage(DenoisingStage):
                                            forward_batch=batch):
                        # Run transformer; follow DMD stage pattern
                        t_expanded_noise = t_cur * torch.ones(
-                            (latent_model_input.shape[0], 1),
+                            (latent_model_input.shape[0], current_num_frames),
                            device=latent_model_input.device,
                            dtype=torch.long)
+
+                        if block_idx == 0 and first_frame_latent is not None:
+                            t_expanded_noise[:, 0] = 0
+
                        pred_noise_btchw = current_model(
                            latent_model_input,
                            prompt_embeds,
                            t_expanded_noise,
-                            kv_cache=self.kv_cache1,
-                            crossattn_cache=self.crossattn_cache,
+                            kv_cache=kv_cache1,
+                            crossattn_cache=crossattn_cache,
                            current_start=(pos_start_base + start_index) *
                            self.frame_seq_length,
                            start_frame=start_index,
@@ -290,12 +348,20 @@ class CausalDMDDenosingStage(DenoisingStage):
                        ).permute(0, 2, 1, 3, 4)

                    # Convert pred noise to pred video with FM Euler scheduler utilities
-                    pred_video_btchw = pred_noise_to_pred_video(
-                        pred_noise=pred_noise_btchw.flatten(0, 1),
-                        noise_input_latent=noise_latents.flatten(0, 1),
-                        timestep=t_expand,
-                        scheduler=self.scheduler).unflatten(
-                            0, pred_noise_btchw.shape[:2])
+                    if boundary_timestep is not None and fastvideo_args.use_add_noise_high and t_cur >= boundary_timestep:
+                        pred_video_btchw = pred_noise_to_x_bound(
+                            pred_noise=pred_noise_btchw.flatten(0, 1),
+                            noise_input_latent=noise_latents.flatten(0, 1),
+                            timestep=t_expanded_noise.flatten(0, 1),
+                            boundary_timestep=torch.ones_like(t_expanded_noise.flatten(0, 1)) * boundary_timestep,
+                            scheduler=self.scheduler).unflatten(0, pred_noise_btchw.shape[:2])
+                    else:
+                        pred_video_btchw = pred_noise_to_pred_video(
+                            pred_noise=pred_noise_btchw.flatten(0, 1),
+                            noise_input_latent=noise_latents.flatten(0, 1),
+                            timestep=t_expanded_noise.flatten(0, 1),
+                            scheduler=self.scheduler).unflatten(
+                                0, pred_noise_btchw.shape[:2])

                    if i < len(timesteps) - 1:
                        next_timestep = timesteps[i + 1] * torch.ones(
@@ -309,17 +375,31 @@ class CausalDMDDenosingStage(DenoisingStage):
                                batch.generator, list) else
                                       batch.generator)).to(self.device)
                        noise_btchw = noise
-                        noise_latents_btchw = self.scheduler.add_noise(
-                            pred_video_btchw.flatten(0, 1),
-                            noise_btchw.flatten(0, 1),
-                            next_timestep).unflatten(0,
-                                                     pred_video_btchw.shape[:2])
+                        if boundary_timestep is not None and fastvideo_args.use_add_noise_high and i < len(high_noise_timesteps) - 1:
+                            noise_latents_btchw = self.scheduler.add_noise_high(
+                                pred_video_btchw.flatten(0, 1),
+                                noise_btchw.flatten(0, 1),
+                                next_timestep,
+                                torch.ones_like(next_timestep) * boundary_timestep).unflatten(0, pred_video_btchw.shape[:2])
+                        elif boundary_timestep is not None and fastvideo_args.use_add_noise_high and i == len(high_noise_timesteps) - 1:
+                            noise_latents_btchw = pred_video_btchw
+                        else:
+                            noise_latents_btchw = self.scheduler.add_noise(
+                                pred_video_btchw.flatten(0, 1),
+                                noise_btchw.flatten(0, 1),
+                                next_timestep).unflatten(0,
+                                                        pred_video_btchw.shape[:2])
                        current_latents = noise_latents_btchw.permute(
                            0, 2, 1, 3, 4)
                    else:
                        current_latents = pred_video_btchw.permute(
                            0, 2, 1, 3, 4)

+                    if block_idx == 0 and first_frame_latent is not None:
+                        noise_latents_btchw = (1 - first_frame_mask) * first_frame_latent + first_frame_mask * noise_latents_btchw.permute(0, 2, 1, 3, 4)
+                        noise_latents_btchw = noise_latents_btchw.permute(0, 2, 1, 3, 4)
+                        current_latents = (1 - first_frame_mask) * first_frame_latent + first_frame_mask * current_latents
+
                    if progress_bar is not None:
                        progress_bar.update()

@@ -327,6 +407,14 @@ class CausalDMDDenosingStage(DenoisingStage):
                latents[:, :, start_index:start_index +
                        current_num_frames, :, :] = current_latents

+                # Decode the latent
+                # first_frame = block_idx == 0
+                # pixels, self.decode_vae_cache = self.decoding_stage.decode(current_latents, fastvideo_args, *self.decode_vae_cache, first_frame=first_frame)
+                # if batch.output is None:
+                #     batch.output = pixels.cpu().float()
+                # else:
+                #     batch.output = torch.cat([batch.output, pixels.cpu().float()], dim=2)
+
                # Re-run with context timestep to update KV cache using clean context
                context_noise = getattr(fastvideo_args.pipeline_config,
                                        "context_noise", 0)
@@ -341,20 +429,57 @@ class CausalDMDDenosingStage(DenoisingStage):
                                        attn_metadata=attn_metadata,
                                        forward_batch=batch):
                    t_expanded_context = t_context.unsqueeze(1)
+
+                    if boundary_timestep is not None:
+                        _ = self.transformer(
+                            context_bcthw,
+                            prompt_embeds,
+                            t_expanded_context,
+                            kv_cache=kv_cache1,
+                            crossattn_cache=crossattn_cache,
+                            current_start=(pos_start_base + start_index) *
+                            self.frame_seq_length,
+                            start_frame=start_index,
+                            **image_kwargs,
+                            **pos_cond_kwargs,
+                        )
+                    
                    _ = current_model(
                        context_bcthw,
                        prompt_embeds,
                        t_expanded_context,
-                        kv_cache=self.kv_cache1,
-                        crossattn_cache=self.crossattn_cache,
+                        kv_cache=kv_cache1,
+                        crossattn_cache=crossattn_cache,
                        current_start=(pos_start_base + start_index) *
                        self.frame_seq_length,
                        start_frame=start_index,
                        **image_kwargs,
                        **pos_cond_kwargs,
                    )
+
                start_index += current_num_frames

+        # remove_indices = torch.arange(3, end=latents.shape[2], step=self.num_frames_per_block)
+        # mask = torch.ones(latents.shape[2], device=latents.device, dtype=torch.bool)
+        # mask[remove_indices] = False
+        # latents = latents[:, :, mask]
+        # latents = latents[:, :, 2:, :, :]
+        # logger.info("latents: %s", latents.shape)
+
+        del kv_cache1
+        del crossattn_cache
+        for module in self.transformer.modules():
+            for name, buf in module._buffers.items():
+                if "k_cache" in name or "v_cache" in name:
+                    module._buffers[name] = None
+        if getattr(self, "transformer_2", None) is not None:
+            for module in self.transformer_2.modules():
+                for name, buf in module._buffers.items():
+                    if "k_cache" in name or "v_cache" in name:
+                        module._buffers[name] = None
+        gc.collect()
+        torch.cuda.empty_cache()
+
        batch.latents = latents
        return batch

@@ -369,6 +494,7 @@ class CausalDMDDenosingStage(DenoisingStage):
            kv_cache_size = self.local_attn_size * self.frame_seq_length
        else:
            kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames
+        # kv_cache_size = 1560 * 19

        for _ in range(self.num_transformer_blocks):
            kv_cache1.append({
@@ -392,7 +518,8 @@ class CausalDMDDenosingStage(DenoisingStage):
                torch.tensor([0], dtype=torch.long, device=device),
            })

-        self.kv_cache1 = kv_cache1
+        # self.kv_cache1 = kv_cache1
+        return kv_cache1

    def _initialize_crossattn_cache(self, batch_size, max_text_len, dtype,
                                    device) -> None:
@@ -421,7 +548,8 @@ class CausalDMDDenosingStage(DenoisingStage):
                "is_init":
                False,
            })
-        self.crossattn_cache = crossattn_cache
+        # self.crossattn_cache = crossattn_cache
+        return crossattn_cache

    def verify_input(self, batch: ForwardBatch,
                     fastvideo_args: FastVideoArgs) -> VerificationResult:
@@ -176,7 +176,7 @@ class DecodingStage(PipelineStage):
            self.maybe_free_model_hooks()

        if fastvideo_args.vae_cpu_offload:
-            self.vae.to("cpu")
+            self.vae = self.vae.to("cpu")

        if torch.backends.mps.is_available():
            del self.vae
@@ -105,27 +105,27 @@ class InputValidationStage(PipelineStage):
                image = load_image(batch.image_path)
            batch.pil_image = image

-        # further processing for ti2v task
-        if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+        # further processing for ti2v task or causal video gen task
+        if (fastvideo_args.pipeline_config.ti2v_task or fastvideo_args.pipeline_config.is_causal) and batch.pil_image is not None:
            img = batch.pil_image
            ih, iw = img.height, img.width
            patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size
            vae_stride = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
            dh, dw = patch_size[1] * vae_stride, patch_size[2] * vae_stride
-            max_area = 704 * 1280
+            max_area = 480 * 832
            ow, oh = best_output_size(iw, ih, dw, dh, max_area)

            scale = max(ow / iw, oh / ih)
            img = img.resize((round(iw * scale), round(ih * scale)),
                             Image.LANCZOS)
-            logger.info("resized img height: %s, img width: %s", img.height,
-                        img.width)

            # center-crop
            x1 = (img.width - ow) // 2
            y1 = (img.height - oh) // 2
            img = img.crop((x1, y1, x1 + ow, y1 + oh))
            assert img.width == ow and img.height == oh
+            logger.info("final processed img height: %s, img width: %s",
+                        img.height, img.width)

            # to tensor
            img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(
@@ -21,14 +21,15 @@ class ModelWrapper(torch.distributed.checkpoint.stateful.Stateful):
        state_dict = get_model_state_dict(
            self.model)  # type: ignore[no-any-return]
        # filter out non-trainable parameters
-        param_requires_grad = set([
-            k for k, v in dict(self.model.named_parameters()).items()
-            if v.requires_grad
-        ])
-        state_dict = {
-            k: v
-            for k, v in state_dict.items() if k in param_requires_grad
-        }
+        # param_requires_grad = set([
+        #     k for k, v in dict(self.model.named_parameters()).items()
+        #     if v.requires_grad
+        # ])
+        # state_dict = {
+        #     k: v
+        #     for k, v in state_dict.items() if k in param_requires_grad
+        # }
+        # logger.info(f"state_dict: {state_dict.keys()}")
        return state_dict  # type: ignore

    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
@@ -7,7 +7,7 @@ import time
 from abc import abstractmethod
 from collections import deque
 from collections.abc import Iterator
-from typing import Any
+from typing import Any, Tuple

 import imageio
 import numpy as np
@@ -30,7 +30,7 @@ from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
 from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
    FlowMatchEulerDiscreteScheduler)
-from fastvideo.models.utils import pred_noise_to_pred_video
+from fastvideo.models.utils import pred_noise_to_pred_video, pred_noise_to_x_bound
 from fastvideo.pipelines import (ComposedPipelineBase, ForwardBatch,
                                 TrainingBatch)
 from fastvideo.training.activation_checkpoint import (
@@ -39,7 +39,7 @@ from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.training_utils import (
    EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases,
    get_scheduler, load_distillation_checkpoint, save_distillation_checkpoint,
-    shift_timestep)
+    shift_timestep, unshift_timestep)
 from fastvideo.utils import (is_vsa_available, maybe_download_model,
                             set_random_seed, verify_model_config_and_directory)

@@ -69,6 +69,7 @@ class DistillationPipeline(TrainingPipeline):
    video_latent_shape: tuple[int, ...]
    video_latent_shape_sp: tuple[int, ...]
    train_fake_score_transformer_2: bool = False
+    i2v_prompt_image_pairs: list[dict[str, str]] = []

    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
        raise RuntimeError(
@@ -78,6 +79,10 @@ class DistillationPipeline(TrainingPipeline):
        """Initialize the distillation training pipeline with multiple models."""
        logger.info("Initializing distillation pipeline...")

+        json_path = "data/mixkit_i2v_full.json"
+        with open(json_path, 'r') as f:
+            self.i2v_prompt_image_pairs = json.load(f)['data']
+
        super().initialize_training_pipeline(training_args)

        self.noise_scheduler = self.get_module("scheduler")
@@ -88,10 +93,14 @@ class DistillationPipeline(TrainingPipeline):
        self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
            shift=self.timestep_shift)

-        if self.training_args.boundary_ratio is not None:
-            self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps
+        if self.training_args.override_boundary_ratio is not None:
+            self.boundary_timestep = self.training_args.override_boundary_ratio * self.noise_scheduler.num_train_timesteps
+            logger.info("Overriding boundary timestep to %s", self.boundary_timestep)
        else:
-            self.boundary_timestep = None
+            if self.training_args.boundary_ratio is not None:
+                self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps
+            else:
+                self.boundary_timestep = None

        if training_args.real_score_model_path:
            logger.info("Loading real score transformer from: %s",
@@ -260,24 +269,25 @@ class DistillationPipeline(TrainingPipeline):
                                self.num_train_timestep)

        self.real_score_guidance_scale = self.training_args.real_score_guidance_scale
+        self.real_score_guidance_scale_2 = self.training_args.real_score_guidance_scale_2

-        self.generator_ema: EMA_FSDP | None = None
-        self.generator_ema_2: EMA_FSDP | None = None
-        if (self.training_args.ema_decay
-                is not None) and (self.training_args.ema_decay > 0.0):
-            self.generator_ema = EMA_FSDP(self.transformer,
-                                          decay=self.training_args.ema_decay)
-            logger.info("Initialized generator EMA with decay=%s",
-                        self.training_args.ema_decay)
+        # self.generator_ema: EMA_FSDP | None = None
+        # self.generator_ema_2: EMA_FSDP | None = None
+        # if (self.training_args.ema_decay
+        #         is not None) and (self.training_args.ema_decay > 0.0):
+        #     self.generator_ema = EMA_FSDP(self.transformer,
+        #                                   decay=self.training_args.ema_decay)
+        #     logger.info("Initialized generator EMA with decay=%s",
+        #                 self.training_args.ema_decay)

-            # Initialize EMA for transformer_2 if it exists
-            if self.transformer_2 is not None:
-                self.generator_ema_2 = EMA_FSDP(
-                    self.transformer_2, decay=self.training_args.ema_decay)
-                logger.info("Initialized generator EMA_2 with decay=%s",
-                            self.training_args.ema_decay)
-        else:
-            logger.info("Generator EMA disabled (ema_decay <= 0.0)")
+        #     # Initialize EMA for transformer_2 if it exists
+        #     if self.transformer_2 is not None:
+        #         self.generator_ema_2 = EMA_FSDP(
+        #             self.transformer_2, decay=self.training_args.ema_decay)
+        #         logger.info("Initialized generator EMA_2 with decay=%s",
+        #                     self.training_args.ema_decay)
+        # else:
+        #     logger.info("Generator EMA disabled (ema_decay <= 0.0)")

    def load_module_from_path(self, model_path: str, module_type: str,
                              training_args: "TrainingArgs"):
@@ -541,6 +551,15 @@ class DistillationPipeline(TrainingPipeline):
            self.generator_ema_2.decay = original_decay_2
            logger.info("EMA_2 reset completed")

+    def _get_real_score_guidance_scale(self, timestep: torch.Tensor):
+        """
+        Get the appropriate real score guidance scale based on timestep and boundary logic.
+        """
+        if self.real_score_transformer_2 is not None and self.boundary_timestep is not None:
+            if timestep.item() < self.boundary_timestep:
+                return self.real_score_guidance_scale_2
+        return self.real_score_guidance_scale
+
    def _get_real_score_transformer(self, timestep: torch.Tensor):
        """
        Get the appropriate real score transformer based on timestep and boundary logic.
@@ -615,7 +634,11 @@ class DistillationPipeline(TrainingPipeline):
            noisy_latent, timestep, training_batch.conditional_dict,
            training_batch)

-        pred_noise = self.transformer(**training_batch.input_kwargs).permute(
+        if self.transformer_2 is not None and self.train_transformer_2:
+            current_transformer = self.transformer_2
+        else:
+            current_transformer = self.transformer
+        pred_noise = current_transformer(**training_batch.input_kwargs).permute(
            0, 2, 1, 3, 4)
        pred_video = pred_noise_to_pred_video(
            pred_noise=pred_noise.flatten(0, 1),
@@ -626,7 +649,9 @@ class DistillationPipeline(TrainingPipeline):
        return pred_video

    def _generator_multi_step_simulation_forward(
-            self, training_batch: TrainingBatch) -> torch.Tensor:
+            self,
+            training_batch: TrainingBatch,
+            exit_flags: list[int] | None = None) -> Tuple[torch.Tensor, float]:
        """Forward pass through student transformer matching inference procedure."""
        latents = training_batch.latents
        dtype = latents.dtype
@@ -662,13 +687,17 @@ class DistillationPipeline(TrainingPipeline):
            with torch.no_grad():
                for step_idx in range(max_target_idx):
                    current_timestep = self.denoising_step_list[step_idx]
+                    if self.transformer_2 is not None and current_timestep < self.boundary_timestep:
+                        current_transformer = self.transformer_2
+                    else:
+                        current_transformer = self.transformer
                    current_timestep_tensor = current_timestep * torch.ones(
                        1, device=self.device, dtype=torch.long)
                    # Run student model to get flow prediction
                    training_batch_temp = self._build_distill_input_kwargs(
                        current_noise_latents, current_timestep_tensor,
                        training_batch.conditional_dict, training_batch)
-                    pred_flow = self.transformer(
+                    pred_flow = current_transformer(
                        **training_batch_temp.input_kwargs).permute(
                            0, 2, 1, 3, 4)
                    pred_clean = pred_noise_to_pred_video(
@@ -711,8 +740,12 @@ class DistillationPipeline(TrainingPipeline):
        training_batch = self._build_distill_input_kwargs(
            noisy_input, target_timestep, training_batch.conditional_dict,
            training_batch)
-        pred_noise = self.transformer(**training_batch.input_kwargs).permute(
-            0, 2, 1, 3, 4)
+        if self.transformer_2 is not None:
+            pred_noise = self.transformer_2(**training_batch.input_kwargs).permute(
+                0, 2, 1, 3, 4)
+        else:
+            pred_noise = self.transformer(**training_batch.input_kwargs).permute(
+                0, 2, 1, 3, 4)
        pred_video = pred_noise_to_pred_video(
            pred_noise=pred_noise.flatten(0, 1),
            noise_input_latent=noisy_input.flatten(0, 1),
@@ -720,28 +753,58 @@ class DistillationPipeline(TrainingPipeline):
            scheduler=self.noise_scheduler).unflatten(0, pred_noise.shape[:2])
        training_batch.dmd_latent_vis_dict[
            "generator_timestep"] = target_timestep.float().detach()
-        return pred_video
+        return pred_video, target_timestep.item()

    def _dmd_forward(self, generator_pred_video: torch.Tensor,
-                     training_batch: TrainingBatch) -> torch.Tensor:
+                     clean_generator_pred_video: torch.Tensor,
+                     training_batch: TrainingBatch,
+                     exit_timestep: float) -> torch.Tensor:
        """Compute DMD (Diffusion Model Distillation) loss."""
        original_latent = generator_pred_video
+        clean_original_latent = clean_generator_pred_video
+        boundary_timestep_pre_shift = int(unshift_timestep(self.boundary_timestep, self.timestep_shift))
        with torch.no_grad():
-            timestep = torch.randint(0,
-                                     self.num_train_timestep, [1],
+            # Detect whether we're doing MoE DMD
+            num_high_noise_timesteps = torch.sum(self.noise_scheduler.timesteps >= self.boundary_timestep).item()
+            if self.boundary_timestep is not None:
+                if exit_timestep < self.boundary_timestep and self.fake_score_transformer_2 is not None:
+                    train_critic_2 = torch.rand(1, generator=self.noise_random_generator).item() >= (num_high_noise_timesteps / self.num_train_timestep)
+                    decision = torch.tensor(1.0 if train_critic_2 else 0.0,
+                                    device=self.device)
+                    world_group = get_world_group()
+                    if world_group.world_size > 1:
+                        world_group.broadcast(decision, src=0)
+                    train_critic_2 = decision.item() == 1.0
+
+                    if train_critic_2:
+                        start_timestep = 0
+                        end_timestep = boundary_timestep_pre_shift - 2
+                    else:
+                        start_timestep = boundary_timestep_pre_shift + 4
+                        end_timestep = self.num_train_timestep
+                elif exit_timestep < self.boundary_timestep and self.fake_score_transformer_2 is None:
+                    start_timestep = 0
+                    end_timestep = self.num_train_timestep
+                else:
+                    start_timestep = boundary_timestep_pre_shift + 4
+                    end_timestep = self.num_train_timestep
+            else:
+                start_timestep = 0
+                end_timestep = self.num_train_timestep
+
+            timestep = torch.randint(start_timestep,
+                                     end_timestep, [1],
                                     device=self.device,
                                     dtype=torch.long)
-            world_group = get_world_group()
-            if world_group.world_size > 1:
-                world_group.broadcast(timestep, src=0)

            timestep = shift_timestep(
                timestep,
-                self.timestep_shift,  # type: ignore
-                self.num_train_timestep)
+                self.timestep_shift)  # type: ignore

            timestep = timestep.clamp(self.min_timestep, self.max_timestep)

+            # logger.info(f"_dmd_forward: rank {self.global_rank} timestep: {timestep}", local_main_process_only=False)
+
            noise = torch.randn(self.video_latent_shape,
                                device=self.device,
                                dtype=generator_pred_video.dtype)
@@ -751,26 +814,63 @@ class DistillationPipeline(TrainingPipeline):
                                  n=self.sp_world_size).contiguous()
                noise = noise[:, self.rank_in_sp_group, :, :, :, :]

-            noisy_latent = self.noise_scheduler.add_noise(
-                generator_pred_video.flatten(0, 1), noise.flatten(0, 1),
-                timestep).detach().unflatten(0,
-                                             (1, generator_pred_video.shape[1]))
+            moe_use_add_noise_high = self.boundary_timestep is not None and self.training_args.use_add_noise_high and exit_timestep >= self.boundary_timestep
+
+            if moe_use_add_noise_high:
+                # if timestep < self.boundary_timestep:
+                #     noisy_latent = self.noise_scheduler.add_noise(
+                #         clean_generator_pred_video.flatten(0, 1), noise.flatten(0, 1),
+                #         timestep).detach().unflatten(0,
+                #                                     (1, clean_original_latent.shape[1]))
+                # else:
+                noisy_latent = self.noise_scheduler.add_noise_high(
+                    generator_pred_video.flatten(0, 1), noise.flatten(0, 1),
+                    timestep,
+                    self.boundary_timestep * torch.ones_like(timestep)).detach().unflatten(0,
+                                                                                        (1, generator_pred_video.shape[1]))
+            else:
+                noisy_latent = self.noise_scheduler.add_noise(
+                    generator_pred_video.flatten(0, 1), noise.flatten(0, 1),
+                    timestep).detach().unflatten(0,
+                                                (1, generator_pred_video.shape[1]))

            # fake_score_transformer forward
            training_batch = self._build_distill_input_kwargs(
                noisy_latent, timestep, training_batch.conditional_dict,
                training_batch)
+            assert int(timestep) != self.boundary_timestep, "Timestep is equal to boundary timestep"
            current_fake_score_transformer = self._get_fake_score_transformer(
                timestep)
            fake_score_pred_noise = current_fake_score_transformer(
                **training_batch.input_kwargs).permute(0, 2, 1, 3, 4)

-            faker_score_pred_video = pred_noise_to_pred_video(
-                pred_noise=fake_score_pred_noise.flatten(0, 1),
-                noise_input_latent=noisy_latent.flatten(0, 1),
-                timestep=timestep,
-                scheduler=self.noise_scheduler).unflatten(
-                    0, fake_score_pred_noise.shape[:2])
+            if moe_use_add_noise_high:
+                # if timestep < self.boundary_timestep:
+                #     raise ValueError("Timestep is less than boundary timestep")
+                #     faker_score_pred_video = pred_noise_to_pred_video(
+                #         pred_noise=fake_score_pred_noise.flatten(0, 1),
+                #         noise_input_latent=noisy_latent.flatten(0, 1),
+                #         timestep=timestep,
+                #         scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
+                # else:
+                faker_score_pred_video = pred_noise_to_x_bound(
+                    pred_noise=fake_score_pred_noise.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    boundary_timestep=torch.ones_like(timestep) * self.boundary_timestep,
+                    scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
+                clean_faker_score_pred_video = pred_noise_to_pred_video(
+                    pred_noise=fake_score_pred_noise.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
+            else:
+                faker_score_pred_video = pred_noise_to_pred_video(
+                    pred_noise=fake_score_pred_noise.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    scheduler=self.noise_scheduler).unflatten(
+                        0, fake_score_pred_noise.shape[:2])

            # real_score_transformer cond forward
            training_batch = self._build_distill_input_kwargs(
@@ -781,12 +881,33 @@ class DistillationPipeline(TrainingPipeline):
            real_score_pred_noise_cond = current_real_score_transformer(
                **training_batch.input_kwargs).permute(0, 2, 1, 3, 4)

-            pred_real_video_cond = pred_noise_to_pred_video(
-                pred_noise=real_score_pred_noise_cond.flatten(0, 1),
-                noise_input_latent=noisy_latent.flatten(0, 1),
-                timestep=timestep,
-                scheduler=self.noise_scheduler).unflatten(
-                    0, real_score_pred_noise_cond.shape[:2])
+            if moe_use_add_noise_high:
+                # if timestep < self.boundary_timestep:
+                #     pred_real_video_cond = pred_noise_to_pred_video(
+                #         pred_noise=real_score_pred_noise_cond.flatten(0, 1),
+                #         noise_input_latent=noisy_latent.flatten(0, 1),
+                #         timestep=timestep,
+                #         scheduler=self.noise_scheduler).unflatten(0, real_score_pred_noise_cond.shape[:2])
+                # else:
+                pred_real_video_cond = pred_noise_to_x_bound(
+                    pred_noise=real_score_pred_noise_cond.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    boundary_timestep=torch.ones_like(timestep) * self.boundary_timestep,
+                    scheduler=self.noise_scheduler).unflatten(0, real_score_pred_noise_cond.shape[:2])
+                clean_pred_real_video_cond = pred_noise_to_pred_video(
+                    pred_noise=real_score_pred_noise_cond.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    scheduler=self.noise_scheduler).unflatten(
+                        0, real_score_pred_noise_cond.shape[:2])
+            else:
+                pred_real_video_cond = pred_noise_to_pred_video(
+                    pred_noise=real_score_pred_noise_cond.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    scheduler=self.noise_scheduler).unflatten(
+                        0, real_score_pred_noise_cond.shape[:2])

            # real_score_transformer uncond forward
            training_batch = self._build_distill_input_kwargs(
@@ -796,16 +917,41 @@ class DistillationPipeline(TrainingPipeline):
            real_score_pred_noise_uncond = current_real_score_transformer(
                **training_batch.input_kwargs).permute(0, 2, 1, 3, 4)

-            pred_real_video_uncond = pred_noise_to_pred_video(
-                pred_noise=real_score_pred_noise_uncond.flatten(0, 1),
-                noise_input_latent=noisy_latent.flatten(0, 1),
-                timestep=timestep,
-                scheduler=self.noise_scheduler).unflatten(
-                    0, real_score_pred_noise_uncond.shape[:2])
+            if moe_use_add_noise_high:
+                # if timestep < self.boundary_timestep:
+                #     pred_real_video_uncond = pred_noise_to_pred_video(
+                #         pred_noise=real_score_pred_noise_uncond.flatten(0, 1),
+                #         noise_input_latent=noisy_latent.flatten(0, 1),
+                #         timestep=timestep,
+                #         scheduler=self.noise_scheduler).unflatten(0, real_score_pred_noise_uncond.shape[:2])
+                # else:
+                pred_real_video_uncond = pred_noise_to_x_bound(
+                    pred_noise=real_score_pred_noise_uncond.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    boundary_timestep=torch.ones_like(timestep) * self.boundary_timestep,
+                    scheduler=self.noise_scheduler).unflatten(0, real_score_pred_noise_uncond.shape[:2])
+                clean_pred_real_video_uncond = pred_noise_to_pred_video(
+                    pred_noise=real_score_pred_noise_uncond.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    scheduler=self.noise_scheduler).unflatten(
+                        0, real_score_pred_noise_uncond.shape[:2])
+            else:
+                pred_real_video_uncond = pred_noise_to_pred_video(
+                    pred_noise=real_score_pred_noise_uncond.flatten(0, 1),
+                    noise_input_latent=noisy_latent.flatten(0, 1),
+                    timestep=timestep,
+                    scheduler=self.noise_scheduler).unflatten(
+                        0, real_score_pred_noise_uncond.shape[:2])

+            if moe_use_add_noise_high:
+                clean_real_score_pred_video = clean_pred_real_video_cond + (
+                    clean_pred_real_video_cond -
+                    clean_pred_real_video_uncond) * self._get_real_score_guidance_scale(timestep)
            real_score_pred_video = pred_real_video_cond + (
                pred_real_video_cond -
-                pred_real_video_uncond) * self.real_score_guidance_scale
+                pred_real_video_uncond) * self._get_real_score_guidance_scale(timestep)

            grad = (faker_score_pred_video - real_score_pred_video) / torch.abs(
                original_latent - real_score_pred_video).mean()
@@ -815,18 +961,32 @@ class DistillationPipeline(TrainingPipeline):
            original_latent.float(),
            (original_latent.float() - grad.float()).detach())

-        training_batch.dmd_latent_vis_dict.update({
-            "training_batch_dmd_fwd_clean_latent":
-            training_batch.latents,
-            "generator_pred_video":
-            original_latent.detach(),
-            "real_score_pred_video":
-            real_score_pred_video.detach(),
-            "faker_score_pred_video":
-            faker_score_pred_video.detach(),
-            "dmd_timestep":
-            timestep.detach(),
-        })
+        if moe_use_add_noise_high:
+            training_batch.dmd_latent_vis_dict.update({
+                "training_batch_dmd_fwd_clean_latent":
+                training_batch.latents,
+                "generator_pred_video":
+                clean_original_latent.detach(),
+                "real_score_pred_video":
+                clean_real_score_pred_video.detach(),
+                "faker_score_pred_video":
+                clean_faker_score_pred_video.detach(),
+                "dmd_timestep":
+                timestep.detach(),
+            })
+        else:
+            training_batch.dmd_latent_vis_dict.update({
+                "training_batch_dmd_fwd_clean_latent":
+                training_batch.latents,
+                "generator_pred_video":
+                original_latent.detach(),
+                "real_score_pred_video":
+                real_score_pred_video.detach(),
+                "faker_score_pred_video":
+                faker_score_pred_video.detach(),
+                "dmd_timestep":
+                timestep.detach(),
+            })

        return dmd_loss

@@ -837,27 +997,55 @@ class DistillationPipeline(TrainingPipeline):
                current_timestep=training_batch.timesteps,
                attn_metadata=training_batch.attn_metadata_vsa):
            if self.training_args.simulate_generator_forward:
-                generator_pred_video = self._generator_multi_step_simulation_forward(
+                generator_pred_video, clean_generator_pred_video, exit_timestep = self._generator_multi_step_simulation_forward(
                    training_batch)
            else:
                generator_pred_video = self._generator_forward(training_batch)

-        fake_score_timestep = torch.randint(0,
-                                            self.num_train_timestep, [1],
+        num_high_noise_timesteps = torch.sum(self.noise_scheduler.timesteps >= self.boundary_timestep).item()
+        # Detect whether we're doing MoE DMD
+        # If yes, then we need to use the appropriate timestep range
+        boundary_timestep_pre_shift = int(unshift_timestep(self.boundary_timestep, self.timestep_shift))
+        if self.boundary_timestep is not None:
+            if exit_timestep < self.boundary_timestep and self.fake_score_transformer_2 is not None:
+                train_critic_2 = torch.rand(1, generator=self.noise_random_generator).item() >= (num_high_noise_timesteps / self.num_train_timestep)
+                decision = torch.tensor(1.0 if train_critic_2 else 0.0,
+                                device=self.device)
+                world_group = get_world_group()
+                if world_group.world_size > 1:
+                    world_group.broadcast(decision, src=0)
+                train_critic_2 = decision.item() == 1.0
+
+                if train_critic_2:
+                    start_timestep = 0
+                    end_timestep = boundary_timestep_pre_shift - 2
+                else:
+                    start_timestep = boundary_timestep_pre_shift + 4
+                    end_timestep = self.num_train_timestep
+            elif exit_timestep < self.boundary_timestep and self.fake_score_transformer_2 is None:
+                start_timestep = 0
+                end_timestep = self.num_train_timestep
+            else:
+                start_timestep = boundary_timestep_pre_shift + 4
+                end_timestep = self.num_train_timestep
+        else:
+            start_timestep = 0
+            end_timestep = self.num_train_timestep
+
+        fake_score_timestep = torch.randint(start_timestep,
+                                            end_timestep, [1],
                                            device=self.device,
                                            dtype=torch.long)
-        world_group = get_world_group()
-        if world_group.world_size > 1:
-            world_group.broadcast(fake_score_timestep, src=0)

        fake_score_timestep = shift_timestep(
            fake_score_timestep,
-            self.timestep_shift,  # type: ignore
-            self.num_train_timestep)
+            self.timestep_shift)  # type: ignore

        fake_score_timestep = fake_score_timestep.clamp(self.min_timestep,
                                                        self.max_timestep)

+        # logger.info(f"faker_score_forward: rank {self.global_rank} fake_score_timestep: {fake_score_timestep}", local_main_process_only=False)
+
        fake_score_noise = torch.randn(self.video_latent_shape,
                                       device=self.device,
                                       dtype=generator_pred_video.dtype)
@@ -867,11 +1055,26 @@ class DistillationPipeline(TrainingPipeline):
                                         n=self.sp_world_size).contiguous()
            fake_score_noise = fake_score_noise[:, self.
                                                rank_in_sp_group, :, :, :, :]
+        
+        moe_use_add_noise_high = self.boundary_timestep is not None and self.training_args.use_add_noise_high and exit_timestep >= self.boundary_timestep

-        noisy_generator_pred_video = self.noise_scheduler.add_noise(
-            generator_pred_video.flatten(0, 1), fake_score_noise.flatten(0, 1),
-            fake_score_timestep).unflatten(0,
-                                           (1, generator_pred_video.shape[1]))
+        if moe_use_add_noise_high:
+            # if fake_score_timestep < self.boundary_timestep:
+            #     noisy_generator_pred_video = self.noise_scheduler.add_noise(
+            #         clean_generator_pred_video.flatten(0, 1), fake_score_noise.flatten(0, 1),
+            #         fake_score_timestep).unflatten(0,
+            #                                     (1, clean_generator_pred_video.shape[1]))
+            # else:
+            noisy_generator_pred_video = self.noise_scheduler.add_noise_high(
+                generator_pred_video.flatten(0, 1), fake_score_noise.flatten(0, 1),
+                fake_score_timestep,
+                self.boundary_timestep * torch.ones_like(fake_score_timestep)).unflatten(0,
+                                                                                    (1, generator_pred_video.shape[1]))
+        else:
+            noisy_generator_pred_video = self.noise_scheduler.add_noise(
+                generator_pred_video.flatten(0, 1), fake_score_noise.flatten(0, 1),
+                fake_score_timestep).unflatten(0,
+                                            (1, generator_pred_video.shape[1]))

        with set_forward_context(current_timestep=training_batch.timesteps,
                                 attn_metadata=training_batch.attn_metadata):
@@ -879,25 +1082,60 @@ class DistillationPipeline(TrainingPipeline):
                noisy_generator_pred_video, fake_score_timestep,
                training_batch.conditional_dict, training_batch)

+            assert int(fake_score_timestep) != self.boundary_timestep, "Timestep is equal to boundary timestep"
            current_fake_score_transformer = self._get_fake_score_transformer(
                fake_score_timestep)
            fake_score_pred_noise = current_fake_score_transformer(
                **training_batch.input_kwargs).permute(0, 2, 1, 3, 4)

-        target = fake_score_noise - generator_pred_video
-        flow_matching_loss = torch.mean((fake_score_pred_noise - target)**2)
+        if moe_use_add_noise_high:
+            # if fake_score_timestep < self.boundary_timestep:
+            #     pred_video = pred_noise_to_pred_video(
+            #         pred_noise=fake_score_pred_noise.flatten(0, 1),
+            #         noise_input_latent=noisy_generator_pred_video.flatten(0, 1),
+            #         timestep=fake_score_timestep,
+            #         scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
+            #     flow_matching_loss = torch.mean((pred_video - clean_generator_pred_video)**2)
+            # else:
+            pred_video = pred_noise_to_x_bound(
+                pred_noise=fake_score_pred_noise.flatten(0, 1),
+                noise_input_latent=noisy_generator_pred_video.flatten(0, 1),
+                timestep=fake_score_timestep,
+                boundary_timestep=torch.ones_like(fake_score_timestep) * self.boundary_timestep,
+                scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
+            flow_matching_loss = torch.mean((pred_video - generator_pred_video)**2)
+        # else:
+        elif self.boundary_timestep is not None and self.training_args.use_add_noise_high and exit_timestep < self.boundary_timestep:
+            pred_video = pred_noise_to_pred_video(
+                pred_noise=fake_score_pred_noise.flatten(0, 1),
+                noise_input_latent=noisy_generator_pred_video.flatten(0, 1),
+                timestep=fake_score_timestep,
+                scheduler=self.noise_scheduler).unflatten(0, fake_score_pred_noise.shape[:2])
+            flow_matching_loss = torch.mean((pred_video - generator_pred_video)**2)
+        else:
+            # This is mathematically equivalent to pred_video - generator_pred_video for the low noise case
+            target = fake_score_noise - generator_pred_video
+            flow_matching_loss = torch.mean((fake_score_pred_noise - target)**2)

-        training_batch.fake_score_latent_vis_dict = {
-            "training_batch_fakerscore_fwd_clean_latent":
-            training_batch.latents,
-            "generator_pred_video": generator_pred_video,
-            "fake_score_timestep": fake_score_timestep,
-        }
+        if moe_use_add_noise_high:
+            training_batch.fake_score_latent_vis_dict = {
+                "training_batch_fakerscore_fwd_clean_latent":
+                training_batch.latents,
+                "generator_pred_video": clean_generator_pred_video,
+                "fake_score_timestep": fake_score_timestep,
+            }
+        else:
+            training_batch.fake_score_latent_vis_dict = {
+                "training_batch_fakerscore_fwd_clean_latent":
+                training_batch.latents,
+                "generator_pred_video": generator_pred_video,
+                "fake_score_timestep": fake_score_timestep,
+            }

        return training_batch, flow_matching_loss

    def _clip_model_grad_norm_(self, training_batch: TrainingBatch,
-                               transformer) -> TrainingBatch:
+                               transformer) -> float:

        max_grad_norm = self.training_args.max_grad_norm

@@ -913,12 +1151,11 @@ class DistillationPipeline(TrainingPipeline):
            grad_norm = grad_norm.item() if grad_norm is not None else 0.0
        else:
            grad_norm = 0.0
-        training_batch.grad_norm = grad_norm
-        return training_batch
+        return grad_norm

    def _prepare_dit_inputs(self,
                            training_batch: TrainingBatch) -> TrainingBatch:
-        super()._prepare_dit_inputs(training_batch)
+        # super()._prepare_dit_inputs(training_batch, prepare_timesteps=prepare_timesteps)
        conditional_dict = {
            "encoder_hidden_states": training_batch.encoder_hidden_states,
            "encoder_attention_mask": training_batch.encoder_attention_mask,
@@ -929,10 +1166,17 @@ class DistillationPipeline(TrainingPipeline):
                "encoder_attention_mask": self.negative_prompt_attention_mask,
            }
            training_batch.unconditional_dict = unconditional_dict
+        else:
+            unconditional_dict = {
+                "encoder_hidden_states": training_batch.encoder_hidden_states,
+                "encoder_attention_mask": training_batch.encoder_attention_mask,
+            }
+            training_batch.unconditional_dict = unconditional_dict

        training_batch.dmd_latent_vis_dict = {}
        training_batch.fake_score_latent_vis_dict = {}

+        training_batch.timesteps = 0
        training_batch.conditional_dict = conditional_dict
        training_batch.raw_latent_shape = training_batch.latents.shape
        training_batch.latents = training_batch.latents.permute(0, 2, 1, 3, 4)
@@ -967,6 +1211,8 @@ class DistillationPipeline(TrainingPipeline):
            batches.append(batch)

        self.optimizer.zero_grad()
+        if self.transformer_2 is not None:
+            self.optimizer_2.zero_grad()
        total_dmd_loss = 0.0
        dmd_latent_vis_dict = {}
        fake_score_latent_vis_dict = {}
@@ -978,7 +1224,7 @@ class DistillationPipeline(TrainingPipeline):
                        current_timestep=batch_gen.timesteps,
                        attn_metadata=batch_gen.attn_metadata_vsa):
                    if self.training_args.simulate_generator_forward:
-                        generator_pred_video = self._generator_multi_step_simulation_forward(
+                        generator_pred_video, exit_timestep = self._generator_multi_step_simulation_forward(
                            batch_gen)
                    else:
                        generator_pred_video = self._generator_forward(
@@ -988,7 +1234,8 @@ class DistillationPipeline(TrainingPipeline):
                                         attn_metadata=batch_gen.attn_metadata):
                    dmd_loss = self._dmd_forward(
                        generator_pred_video=generator_pred_video,
-                        training_batch=batch_gen)
+                        training_batch=batch_gen,
+                        exit_timestep=exit_timestep)

                with set_forward_context(
                        current_timestep=batch_gen.timesteps,
@@ -997,12 +1244,20 @@ class DistillationPipeline(TrainingPipeline):
                total_dmd_loss += dmd_loss.detach().item()

            # Only clip gradients for the model that is currently training
-            self._clip_model_grad_norm_(batch_gen, self.transformer)
-            for param in self.transformer.parameters():
-                # check if the gradient is not None and not zero
-                assert param.grad is not None and param.grad.abs().sum() > 0
-            self.optimizer.step()
-            self.optimizer.zero_grad(set_to_none=True)
+            if self.transformer_2 is not None and self.train_transformer_2:
+                self._clip_model_grad_norm_(batch_gen, self.transformer_2)
+                for param in self.transformer_2.parameters():
+                    # check if the gradient is not None and not zero
+                    assert param.grad is not None and param.grad.abs().sum() > 0
+                self.optimizer_2.step()
+                self.optimizer_2.zero_grad(set_to_none=True)
+            else:
+                self._clip_model_grad_norm_(batch_gen, self.transformer)
+                for param in self.transformer.parameters():
+                    # check if the gradient is not None and not zero
+                    assert param.grad is not None and param.grad.abs().sum() > 0
+                self.optimizer.step()
+                self.optimizer.zero_grad(set_to_none=True)

            if self.generator_ema is not None:
                self.generator_ema.update(self.transformer)
@@ -1161,18 +1416,17 @@ class DistillationPipeline(TrainingPipeline):
        else:
            logger.info("  Generator EMA disabled")

-        if self.generator_ema is not None:
-            logger.info("  Generator EMA enabled with decay: %s",
+        if getattr(self, 'generator_ema_2', None) is not None:
+            logger.info("  Generator EMA_2 enabled with decay: %s",
                        self.training_args.ema_decay)
-            logger.info("  Generator EMA start step: %s",
+            logger.info("  Generator EMA_2 start step: %s",
                        self.training_args.ema_start_step)
        else:
-            logger.info("  Generator EMA disabled")
+            logger.info("  Generator EMA_2 disabled")

    @torch.no_grad()
-    def _log_validation(self, transformer, training_args, global_step) -> None:
+    def _log_validation(self, training_args, global_step) -> None:
        training_args.inference_mode = True
-        training_args.dit_cpu_offload = True
        if not training_args.log_validation:
            return
        if self.validation_pipeline is None:
@@ -1199,7 +1453,7 @@ class DistillationPipeline(TrainingPipeline):
                                           num_workers=0)

        # Set both transformers to eval mode
-        transformer.eval()
+        self.transformer.eval()
        if hasattr(self, 'transformer_2') and self.transformer_2 is not None:
            self.transformer_2.eval()

@@ -1210,23 +1464,19 @@ class DistillationPipeline(TrainingPipeline):
        ema_2_context = None

        if use_ema_for_validation:
-            logger.info("Using EMA model for validation")
            # Use self.transformer for consistency (the passed transformer should be self.transformer anyway)
-            validation_transformer = self.transformer
-            if self.generator_ema is not None:
-                ema_context = self.generator_ema.apply_to_model(
-                    validation_transformer)
+            # if self.generator_ema is not None:
+            ema_context = self.generator_ema.apply_to_model(
+                self.transformer)
+            logger.info("Using EMA model for validation")

            # Handle transformer_2 EMA if available
            if hasattr(
                    self, 'transformer_2'
-            ) and self.transformer_2 is not None and self.generator_ema_2 is not None:
+            ) and self.transformer_2 is not None:
                ema_2_context = self.generator_ema_2.apply_to_model(
                    self.transformer_2)
                logger.info("Using EMA_2 model for transformer_2 validation")
-        else:
-            # Use self.transformer for consistency, but the passed transformer should be the same
-            validation_transformer = self.transformer

        validation_steps = training_args.validation_sampling_steps.split(",")
        validation_steps = [int(step) for step in validation_steps]
@@ -1236,13 +1486,64 @@ class DistillationPipeline(TrainingPipeline):
        num_sp_groups = world_group.world_size // self.sp_group.world_size
        # Process each validation prompt for each validation step
        for num_inference_steps in validation_steps:
-            logger.info("rank: %s: num_inference_steps: %s",
-                        self.global_rank,
-                        num_inference_steps,
-                        local_main_process_only=False)
+            # logger.info("rank: %s: num_inference_steps: %s",
+            #             self.global_rank,
+            #             num_inference_steps,
+            #             local_main_process_only=False)
            step_videos: list[np.ndarray] = []
            step_captions: list[str] = []

+            def run_i2v_validation_with_ema(
+                steps: int) -> tuple[list[np.ndarray], list[str]]:
+                videos: list[np.ndarray] = []
+                captions: list[str] = []
+                for iter in range(len(self.i2v_prompt_image_pairs) // world_group.world_size):
+                    validation_batch = self.i2v_prompt_image_pairs[iter * world_group.world_size + self.global_rank]
+                    batch = self._prepare_validation_batch(
+                        sampling_param, training_args, validation_batch, steps)
+
+                    negative_prompt = batch.negative_prompt
+                    batch_negative = ForwardBatch(
+                        data_type="video",
+                        prompt=negative_prompt,
+                        prompt_embeds=[],
+                        prompt_attention_mask=[],
+                    )
+                    result_batch = self.validation_pipeline.prompt_encoding_stage(  # type: ignore
+                        batch_negative, training_args)
+                    self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
+                        0], result_batch.prompt_attention_mask[0]
+
+                    logger.info(
+                        "rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
+                        self.global_rank,
+                        self.rank_in_sp_group,
+                        batch.prompt,
+                        local_main_process_only=False)
+
+                    assert batch.prompt is not None and isinstance(
+                        batch.prompt, str)
+                    captions.append(batch.prompt)
+
+                    # Run validation inference
+                    with torch.no_grad():
+                        output_batch = self.validation_pipeline.forward(
+                            batch, training_args)
+                    samples = output_batch.output
+                    if self.rank_in_sp_group != 0:
+                        return videos, captions
+
+                    # Process outputs
+                    video = rearrange(samples, "b c t h w -> t b c h w")
+                    frames = []
+                    for x in video:
+                        x = torchvision.utils.make_grid(x, nrow=6)
+                        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+                        frames.append((x * 255).numpy().astype(np.uint8))
+                    videos.append(frames)
+
+                return videos, captions
+
            # Helper function to run validation with optional EMA contexts
            def run_validation_with_ema(
                    steps: int) -> tuple[list[np.ndarray], list[str]]:
@@ -1311,6 +1612,22 @@ class DistillationPipeline(TrainingPipeline):
                step_videos, step_captions = run_validation_with_ema(
                    num_inference_steps)

+            if ema_context is not None and ema_2_context is not None:
+                with ema_context, ema_2_context:
+                    i2v_step_videos, i2v_step_captions = run_i2v_validation_with_ema(
+                        num_inference_steps)
+            elif ema_context is not None:
+                with ema_context:
+                    i2v_step_videos, i2v_step_captions = run_i2v_validation_with_ema(
+                        num_inference_steps)
+            elif ema_2_context is not None:
+                with ema_2_context:
+                    i2v_step_videos, i2v_step_captions = run_i2v_validation_with_ema(
+                        num_inference_steps)
+            else:
+                i2v_step_videos, i2v_step_captions = run_i2v_validation_with_ema(
+                    num_inference_steps)
+
            # Log validation results for this step
            world_group = get_world_group()
            num_sp_groups = world_group.world_size // self.sp_group.world_size
@@ -1355,8 +1672,47 @@ class DistillationPipeline(TrainingPipeline):
                    world_group.send_object(step_videos, dst=0)
                    world_group.send_object(step_captions, dst=0)

+                # sp_group leaders wait for global rank 0 to finish logging
+                world_group.barrier()
+
+                if self.global_rank == 0:
+                    all_i2v_videos = i2v_step_videos
+                    all_i2v_captions = i2v_step_captions
+
+                    # Receive from other sp_group leaders
+                    for sp_group_idx in range(1, num_sp_groups):
+                        src_rank = sp_group_idx * self.sp_world_size  # Global rank of other sp_group leaders
+                        recv_videos = world_group.recv_object(src=src_rank)
+                        recv_captions = world_group.recv_object(src=src_rank)
+                        all_i2v_videos.extend(recv_videos)
+                        all_i2v_captions.extend(recv_captions)
+
+                    video_filenames = []
+                    for i, (video, caption) in enumerate(
+                            zip(all_i2v_videos, all_i2v_captions, strict=True)):
+                        os.makedirs(training_args.output_dir, exist_ok=True)
+                        filename = os.path.join(
+                            training_args.output_dir,
+                            f"validation_step_{global_step}_inference_steps_{num_inference_steps}_i2v_video_{i}.mp4"
+                        )
+                        imageio.mimsave(filename, video, fps=sampling_param.fps)
+                        video_filenames.append(filename)
+
+                    logs = {
+                        f"i2v_validation_videos_{num_inference_steps}_steps": [
+                            wandb.Video(filename, caption=caption)
+                            for filename, caption in zip(
+                                video_filenames, all_i2v_captions, strict=True)
+                        ]
+                    }
+                    wandb.log(logs, step=global_step)
+                else:
+                    world_group.send_object(i2v_step_videos, dst=0)
+                    world_group.send_object(i2v_step_captions, dst=0)
+
+
        # Re-enable gradients for training - set both transformers back to train mode
-        transformer.train()
+        self.transformer.train()
        if hasattr(self, 'transformer_2') and self.transformer_2 is not None:
            self.transformer_2.train()
        gc.collect()
@@ -1681,6 +2037,8 @@ class DistillationPipeline(TrainingPipeline):
                if self.training_args.use_ema and self.is_ema_ready():
                    self.save_ema_weights(self.training_args.output_dir, step)

+                self.sp_group.barrier()
+
            if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
                if self.training_args.log_visualization:
                    self.visualize_intermediate_latents(training_batch,
@@ -2,7 +2,8 @@
 import copy
 import time
 from collections import deque
-from typing import Any
+from typing import Any, Tuple
+import gc

 import numpy as np
 import torch
@@ -19,13 +20,15 @@ from fastvideo.fastvideo_args import TrainingArgs
 from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger

+from fastvideo.configs.sample import SamplingParam
 from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import SelfForcingFlowMatchScheduler
-from fastvideo.models.utils import pred_noise_to_pred_video
-from fastvideo.pipelines import TrainingBatch
+from fastvideo.models.utils import pred_noise_to_pred_video, pred_noise_to_x_bound
+from fastvideo.pipelines import TrainingBatch, ForwardBatch
 from fastvideo.training.distillation_pipeline import DistillationPipeline
 from fastvideo.training.training_utils import (EMA_FSDP,
                                               save_distillation_checkpoint)
 from fastvideo.utils import is_vsa_available, set_random_seed
+from copy import deepcopy

 logger = init_logger(__name__)

@@ -47,10 +50,14 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
        """Initialize the self-forcing training pipeline."""
        logger.info("Initializing self-forcing distillation pipeline...")

+        # self.generator_ema: EMA_FSDP | None = None
+        # self.generator_ema_2: EMA_FSDP | None = None
+        
+        super().initialize_training_pipeline(training_args)
+
        self.generator_ema: EMA_FSDP | None = None
        self.generator_ema_2: EMA_FSDP | None = None

-        super().initialize_training_pipeline(training_args)
        try:
            logger.info("RANK: %s, entered initialize_training_pipeline",
                        self.global_rank,
@@ -60,7 +67,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):

        self.noise_scheduler = SelfForcingFlowMatchScheduler(
            num_inference_steps=1000,
-            shift=5.0,
+            shift=self.timestep_shift,
            sigma_min=0.0,
            extra_one_step=True,
            training=True)
@@ -76,8 +83,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
        self.last_step_only = getattr(training_args, 'last_step_only', False)
        self.context_noise = getattr(training_args, 'context_noise', 0)

-        self.kv_cache1: list[dict[str, Any]] | None = None
-        self.crossattn_cache: list[dict[str, Any]] | None = None
+        # self.kv_cache1: list[dict[str, Any]] | None = None
+        # self.crossattn_cache: list[dict[str, Any]] | None = None

        logger.info("Self-forcing generator update ratio: %s",
                    self.dfake_gen_update_ratio)
@@ -88,13 +95,13 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
    def generate_and_sync_list(self, num_blocks: int, num_denoising_steps: int,
                               device: torch.device) -> list[int]:
        """Generate and synchronize random exit flags across distributed processes."""
-        logger.info(
-            "RANK: %s, enter generate_and_sync_list blocks=%s steps=%s device=%s",
-            self.global_rank,
-            num_blocks,
-            num_denoising_steps,
-            str(device),
-            local_main_process_only=False)
+        # logger.info(
+        #     "RANK: %s, enter generate_and_sync_list blocks=%s steps=%s device=%s",
+        #     self.global_rank,
+        #     num_blocks,
+        #     num_denoising_steps,
+        #     str(device),
+        #     local_main_process_only=False)
        rank = dist.get_rank() if dist.is_initialized() else 0

        if rank == 0:
@@ -112,12 +119,12 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
            dist.broadcast(indices,
                           src=0)  # Broadcast the random indices to all ranks
        flags = indices.tolist()
-        logger.info(
-            "RANK: %s, exit generate_and_sync_list flags_len=%s first=%s",
-            self.global_rank,
-            len(flags),
-            flags[0] if len(flags) > 0 else None,
-            local_main_process_only=False)
+        # logger.info(
+        #     "RANK: %s, exit generate_and_sync_list flags_len=%s first=%s",
+        #     self.global_rank,
+        #     len(flags),
+        #     flags[0] if len(flags) > 0 else None,
+        #     local_main_process_only=False)
        return flags

    def generator_loss(
@@ -127,17 +134,42 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
        Compute generator loss using DMD-style approach.
        The generator tries to fool the critic (fake_score_transformer).
        """
+        # Check that all params in self.transformer require grad is True
+        # for param in self.transformer.parameters():
+        #     assert param.requires_grad, "Param %s in transformer requires grad is False" % param.name
+        # if self.transformer_2 is not None:
+        #     for param in self.transformer_2.parameters():
+        #         assert param.requires_grad, "Param %s in transformer_2 requires grad is False" % param.name
+        exit_flags = None
+        # If the generator is MoE, randomly sample an exit timestep, and turn on training for one of the dits
+        if self.boundary_timestep is not None and self.transformer_2 is not None:
+            assert self.same_step_across_blocks, "same_step_across_blocks must be True for MoE generator. Otherwise we might need to train both transformers which will cause OOM"
+            exit_flags = self.generate_and_sync_list(1, len(self.denoising_step_list), training_batch.latents.device)
+            # exit_flags = self.generate_and_sync_list(1, 2, training_batch.latents.device)
+            # exit_flags[0] += 2
+            exit_timestep = self.denoising_step_list[exit_flags[0]].item()
+            # logger.info("Exit timestep in generator_loss(): %s", exit_timestep)
+            if exit_timestep < self.boundary_timestep:
+                self.train_transformer_2 = True
+            else:
+                self.train_transformer_2 = False
+        else: # Default to normal single-dit generator
+            self.train_transformer_2 = False
+
        with set_forward_context(
                current_timestep=training_batch.timesteps,
                attn_metadata=training_batch.attn_metadata_vsa):
-            generator_pred_video = self._generator_multi_step_simulation_forward(
-                training_batch)
+            # If the exit timestep is in the high noise region, the generator_pred_video is x_boundary, otherwise it's x_0
+            generator_pred_video, clean_generator_pred_video, exit_timestep = self._generator_multi_step_simulation_forward(
+                training_batch, exit_flags=exit_flags)

        with set_forward_context(current_timestep=training_batch.timesteps,
                                 attn_metadata=training_batch.attn_metadata):
            dmd_loss = self._dmd_forward(
                generator_pred_video=generator_pred_video,
-                training_batch=training_batch)
+                clean_generator_pred_video=clean_generator_pred_video,
+                training_batch=training_batch,
+                exit_timestep=exit_timestep)

        log_dict = {
            "dmdtrain_gradient_norm": torch.tensor(0.0, device=self.device)
@@ -152,6 +184,12 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
        Compute critic loss using flow matching between noise and generator output.
        The critic learns to predict the flow from noise to the generator's output.
        """
+        # for param in self.fake_score_transformer.parameters():
+        #     assert param.requires_grad, "Param %s in fake_score_transformer requires grad is False" % param.name
+        # if self.fake_score_transformer_2 is not None:
+        #     for param in self.fake_score_transformer_2.parameters():
+        #         assert param.requires_grad, "Param %s in fake_score_transformer_2 requires grad is False" % param.name
+
        updated_batch, flow_matching_loss = self.faker_score_forward(
            training_batch)
        training_batch.fake_score_latent_vis_dict = updated_batch.fake_score_latent_vis_dict
@@ -162,7 +200,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
    def _generator_multi_step_simulation_forward(
            self,
            training_batch: TrainingBatch,
-            return_sim_steps: bool = False) -> torch.Tensor:
+            exit_flags: list[int] | None = None) -> Tuple[torch.Tensor, float]:
        """Forward pass through student transformer matching inference procedure with KV cache management.
        
        This function is adapted from the reference self-forcing implementation's inference_with_trajectory
@@ -173,29 +211,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
        batch_size = latents.shape[0]
        initial_latent = getattr(training_batch, 'image_latent', None)

-        # Dynamic frame generation logic (adapted from _run_generator)
-        num_training_frames = getattr(self.training_args, 'num_latent_t', 21)
-
-        # During training, the number of generated frames should be uniformly sampled from
-        # [21, self.num_training_frames], but still being a multiple of self.num_frame_per_block
-        min_num_frames = 20 if self.independent_first_frame else 21
-        max_num_frames = num_training_frames - 1 if self.independent_first_frame else num_training_frames
-        assert max_num_frames % self.num_frame_per_block == 0
-        assert min_num_frames % self.num_frame_per_block == 0
-        max_num_blocks = max_num_frames // self.num_frame_per_block
-        min_num_blocks = min_num_frames // self.num_frame_per_block
-
-        # Sample number of blocks and sync across processes
-        num_generated_blocks = torch.randint(min_num_blocks,
-                                             max_num_blocks + 1, (1, ),
-                                             device=self.device)
-        if dist.is_initialized():
-            dist.broadcast(num_generated_blocks, src=0)
-        num_generated_blocks = num_generated_blocks.item()
-        num_generated_frames = num_generated_blocks * self.num_frame_per_block
-        if self.independent_first_frame and initial_latent is None:
-            num_generated_frames += 1
-            min_num_frames += 1
+        num_generated_frames = getattr(self.training_args, 'num_latent_t', 21)

        # Create noise with dynamic shape
        if initial_latent is not None:
@@ -218,13 +234,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
        batch_size, num_frames, num_channels, height, width = noise.shape

        # Block size calculation
-        if not self.independent_first_frame or (self.independent_first_frame
-                                                and initial_latent is not None):
-            assert num_frames % self.num_frame_per_block == 0
-            num_blocks = num_frames // self.num_frame_per_block
-        else:
-            assert (num_frames - 1) % self.num_frame_per_block == 0
-            num_blocks = (num_frames - 1) // self.num_frame_per_block

        num_input_frames = initial_latent.shape[
            1] if initial_latent is not None else 0
@@ -233,18 +242,14 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
            [batch_size, num_output_frames, num_channels, height, width],
            device=noise.device,
            dtype=noise.dtype)
-
-        def get_model_device(model):
-            if model is None:
-                return "None"
-            try:
-                return next(model.parameters()).device
-            except (StopIteration, AttributeError):
-                return "Unknown"
+        clean_output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype)

        # Step 1: Initialize KV cache to all zeros
        cache_frames = num_generated_frames + num_input_frames
-        self.kv_cache1, self.crossattn_cache = self._initialize_simulation_caches(
+        kv_cache1, crossattn_cache = self._initialize_simulation_caches(
            batch_size, dtype, self.device, max_num_frames=cache_frames)

        # Step 2: Cache context feature
@@ -268,22 +273,33 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                    timestep=training_batch_temp.input_kwargs['timestep'],
                    encoder_hidden_states_image=training_batch_temp.
                    input_kwargs.get('encoder_hidden_states_image'),
-                    kv_cache=self.kv_cache1,
-                    crossattn_cache=self.crossattn_cache,
+                    kv_cache=kv_cache1,
+                    crossattn_cache=crossattn_cache,
                    current_start=current_start_frame * self.frame_seq_length,
                    start_frame=current_start_frame)
            current_start_frame += 1

        # Step 3: Temporal denoising loop
-        all_num_frames = [self.num_frame_per_block] * num_blocks
+        all_num_frames = [self.num_frame_per_block] * 7
+        # all_num_frames[0] = 1
        if self.independent_first_frame and initial_latent is None:
            all_num_frames = [1] + all_num_frames
        num_denoising_steps = len(self.denoising_step_list)
-        exit_flags = self.generate_and_sync_list(len(all_num_frames),
+
+        if exit_flags is None:
+            exit_flags = self.generate_and_sync_list(len(all_num_frames),
                                                 num_denoising_steps,
                                                 device=noise.device)
+        # if exit_flags is None:
+        #     exit_flags = self.generate_and_sync_list(len(all_num_frames), 2, device=noise.device)
+        #     exit_flags[0] += 2
+        exit_timestep = self.denoising_step_list[exit_flags[0]].item()
+
        start_gradient_frame_index = max(0, num_output_frames - 21)

+        high_noise_timesteps = self.denoising_step_list[self.denoising_step_list >= self.boundary_timestep]
+        assert len(high_noise_timesteps) == len(self.denoising_step_list) // 2, f"Incorrect number of high noise timesteps: {len(high_noise_timesteps)}"
+
        for block_index, current_num_frames in enumerate(all_num_frames):
            noisy_input = noise[:, current_start_frame -
                                num_input_frames:current_start_frame +
@@ -302,14 +318,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):

                if self.boundary_timestep is not None and current_timestep < self.boundary_timestep and self.transformer_2 is not None:
                    current_model = self.transformer_2
-                    self._enable_training(self.transformer_2, self.optimizer_2)
-                    self._disable_training(self.transformer, self.optimizer)
                else:
                    current_model = self.transformer
-                    self._enable_training(self.transformer, self.optimizer)
-                    if self.boundary_timestep is not None and self.transformer_2 is not None:
-                        self._disable_training(self.transformer_2,
-                                               self.optimizer_2)

                if not exit_flag:
                    with torch.no_grad():
@@ -327,30 +337,84 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                            input_kwargs['timestep'],
                            encoder_hidden_states_image=training_batch_temp.
                            input_kwargs.get('encoder_hidden_states_image'),
-                            kv_cache=self.kv_cache1,
-                            crossattn_cache=self.crossattn_cache,
+                            kv_cache=kv_cache1,
+                            crossattn_cache=crossattn_cache,
                            current_start=current_start_frame *
                            self.frame_seq_length,
                            start_frame=current_start_frame).permute(
                                0, 2, 1, 3, 4)

-                        denoised_pred = pred_noise_to_pred_video(
-                            pred_noise=pred_flow.flatten(0, 1),
-                            noise_input_latent=noisy_input.flatten(0, 1),
-                            timestep=timestep,
-                            scheduler=self.noise_scheduler).unflatten(
-                                0, pred_flow.shape[:2])
+                        if self.boundary_timestep is not None and self.training_args.use_add_noise_high and current_timestep >= self.boundary_timestep:
+                            denoised_pred = pred_noise_to_x_bound(
+                                pred_noise=pred_flow.flatten(0, 1),
+                                noise_input_latent=noisy_input.flatten(0, 1),
+                                timestep=timestep,
+                                boundary_timestep=torch.ones_like(timestep).flatten(0, 1) * self.boundary_timestep,
+                                scheduler=self.noise_scheduler).unflatten(
+                                    0, pred_flow.shape[:2])
+                        else:
+                            denoised_pred = pred_noise_to_pred_video(
+                                pred_noise=pred_flow.flatten(0, 1),
+                                noise_input_latent=noisy_input.flatten(0, 1),
+                                timestep=timestep,
+                                scheduler=self.noise_scheduler).unflatten(
+                                    0, pred_flow.shape[:2])

                        next_timestep = self.denoising_step_list[index + 1]
-                        noisy_input = self.noise_scheduler.add_noise(
-                            denoised_pred.flatten(0, 1),
-                            torch.randn_like(denoised_pred.flatten(0, 1)),
-                            next_timestep *
-                            torch.ones([batch_size * current_num_frames],
-                                       device=noise.device,
-                                       dtype=torch.long)).unflatten(
-                                           0, denoised_pred.shape[:2])
+
+                        if self.boundary_timestep is not None and self.training_args.use_add_noise_high and index < len(high_noise_timesteps) - 1:
+                            noisy_input = self.noise_scheduler.add_noise_high(
+                                denoised_pred.flatten(0, 1),
+                                torch.randn_like(denoised_pred.flatten(0, 1)),
+                                next_timestep *
+                                torch.ones([batch_size * current_num_frames],
+                                           device=noise.device,
+                                           dtype=torch.long),
+                                self.boundary_timestep *
+                                torch.ones([batch_size * current_num_frames],
+                                           device=noise.device,
+                                           dtype=torch.long)).unflatten(
+                                    0, denoised_pred.shape[:2])
+                        elif self.boundary_timestep is not None and self.training_args.use_add_noise_high and index == len(high_noise_timesteps) - 1:
+                            assert exit_timestep < self.boundary_timestep, "exit timestep is greater than boundary timestep"
+                            noisy_input = denoised_pred
+
+                            # Step 3.3: rerun with boundary timestep to update the high noise's cache
+                            # context_timestep = torch.ones_like(timestep) * self.boundary_timestep
+                            # context_denoised_pred = self.noise_scheduler.add_noise_high(
+                            #     denoised_pred.flatten(0, 1),
+                            #     torch.randn_like(denoised_pred.flatten(0, 1)),
+                            #     context_timestep,
+                            #     self.boundary_timestep * torch.ones_like(context_timestep)).unflatten(0, denoised_pred.shape[:2])
+
+                            # training_batch_bound = self._build_distill_input_kwargs(
+                            #     context_denoised_pred, context_timestep,
+                            #     training_batch.conditional_dict, training_batch)
+
+                            # current_model(
+                            #     hidden_states=training_batch_bound.
+                            #     input_kwargs['hidden_states'],
+                            #     encoder_hidden_states=training_batch_bound.
+                            #     input_kwargs['encoder_hidden_states'],
+                            #     timestep=training_batch_bound.input_kwargs['timestep'],
+                            #     encoder_hidden_states_image=training_batch_bound.
+                            #     input_kwargs.get('encoder_hidden_states_image'),
+                            #     kv_cache=kv_cache1,
+                            #     crossattn_cache=crossattn_cache,
+                            #     current_start=current_start_frame * self.frame_seq_length,
+                            #     start_frame=current_start_frame)
+
+                        else:
+                            noisy_input = self.noise_scheduler.add_noise(
+                                denoised_pred.flatten(0, 1),
+                                torch.randn_like(denoised_pred.flatten(0, 1)),
+                                next_timestep *
+                                torch.ones([batch_size * current_num_frames],
+                                        device=noise.device,
+                                        dtype=torch.long)).unflatten(
+                                            0, denoised_pred.shape[:2])
                else:
+                    # logger.info("Exit timestep in _generator_multi_step_simulation_forward(): %s", current_timestep)
                    # Final prediction with gradient control
                    if current_start_frame < start_gradient_frame_index:
                        with torch.no_grad():
@@ -367,8 +431,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                                input_kwargs['timestep'],
                                encoder_hidden_states_image=training_batch_temp.
                                input_kwargs.get('encoder_hidden_states_image'),
-                                kv_cache=self.kv_cache1,
-                                crossattn_cache=self.crossattn_cache,
+                                kv_cache=kv_cache1,
+                                crossattn_cache=crossattn_cache,
                                current_start=current_start_frame *
                                self.frame_seq_length,
                                start_frame=current_start_frame).permute(
@@ -387,39 +451,98 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                            input_kwargs['timestep'],
                            encoder_hidden_states_image=training_batch_temp.
                            input_kwargs.get('encoder_hidden_states_image'),
-                            kv_cache=self.kv_cache1,
-                            crossattn_cache=self.crossattn_cache,
+                            kv_cache=kv_cache1,
+                            crossattn_cache=crossattn_cache,
                            current_start=current_start_frame *
                            self.frame_seq_length,
                            start_frame=current_start_frame).permute(
                                0, 2, 1, 3, 4)

-                    denoised_pred = pred_noise_to_pred_video(
-                        pred_noise=pred_flow.flatten(0, 1),
-                        noise_input_latent=noisy_input.flatten(0, 1),
-                        timestep=timestep,
-                        scheduler=self.noise_scheduler).unflatten(
-                            0, pred_flow.shape[:2])
+                    if self.boundary_timestep is not None and self.training_args.use_add_noise_high and current_timestep >= self.boundary_timestep:
+                        denoised_pred = pred_noise_to_x_bound(
+                            pred_noise=pred_flow.flatten(0, 1),
+                            noise_input_latent=noisy_input.flatten(0, 1),
+                            timestep=timestep,
+                            boundary_timestep=torch.ones_like(timestep).flatten(0, 1) * self.boundary_timestep,
+                            scheduler=self.noise_scheduler).unflatten(0, pred_flow.shape[:2])
+                        clean_denoised_pred = pred_noise_to_pred_video(
+                            pred_noise=pred_flow.flatten(0, 1),
+                            noise_input_latent=noisy_input.flatten(0, 1),
+                            timestep=timestep,
+                            scheduler=self.noise_scheduler).unflatten(
+                                0, pred_flow.shape[:2])
+                    else:
+                        denoised_pred = pred_noise_to_pred_video(
+                            pred_noise=pred_flow.flatten(0, 1),
+                            noise_input_latent=noisy_input.flatten(0, 1),
+                            timestep=timestep,
+                            scheduler=self.noise_scheduler).unflatten(
+                                0, pred_flow.shape[:2])
+                        clean_denoised_pred = None
                    break

+            assert exit_timestep == current_timestep, "exit timestep is not equal to current timestep"
+
            # Step 3.2: record the model's output
            output[:, current_start_frame:current_start_frame +
                   current_num_frames] = denoised_pred
+            if clean_denoised_pred is not None:
+                clean_output[:, current_start_frame:current_start_frame +
+                    current_num_frames] = clean_denoised_pred

            # Step 3.3: rerun with timestep zero to update the cache
-            context_timestep = torch.ones_like(timestep) * self.context_noise
-            denoised_pred = self.noise_scheduler.add_noise(
-                denoised_pred.flatten(0, 1),
-                torch.randn_like(denoised_pred.flatten(0, 1)),
-                context_timestep).unflatten(0, denoised_pred.shape[:2])
+            if self.boundary_timestep is not None:
+                if exit_timestep < self.boundary_timestep:
+                    context_timestep = torch.ones_like(timestep) * self.context_noise
+                    current_model = self.transformer_2
+                else:
+                    # For high noise expert, the context timestep should be the boundary timestep
+                    # context_timestep = torch.ones_like(timestep) * self.boundary_timestep
+                    context_timestep = torch.ones_like(timestep) * self.context_noise
+                    current_model = self.transformer
+            else:
+                context_timestep = torch.ones_like(timestep) * self.context_noise
+                current_model = self.transformer
+
+            if self.boundary_timestep is not None and self.training_args.use_add_noise_high and current_timestep >= self.boundary_timestep:
+                # denoised_pred = self.noise_scheduler.add_noise_high(
+                #     denoised_pred.flatten(0, 1),
+                #     torch.randn_like(denoised_pred.flatten(0, 1)),
+                #     context_timestep,
+                #     self.boundary_timestep * torch.ones_like(context_timestep)).unflatten(0, denoised_pred.shape[:2])
+                denoised_pred = self.noise_scheduler.add_noise(
+                    clean_denoised_pred.flatten(0, 1),
+                    torch.randn_like(clean_denoised_pred.flatten(0, 1)),
+                    context_timestep).unflatten(0, clean_denoised_pred.shape[:2])
+            else:
+                denoised_pred = self.noise_scheduler.add_noise(
+                    denoised_pred.flatten(0, 1),
+                    torch.randn_like(denoised_pred.flatten(0, 1)),
+                    context_timestep).unflatten(0, denoised_pred.shape[:2])

            with torch.no_grad():
                training_batch_temp = self._build_distill_input_kwargs(
                    denoised_pred, context_timestep,
                    training_batch.conditional_dict, training_batch)

+                if self.boundary_timestep is not None and exit_timestep < self.boundary_timestep:
+                    self.transformer(
+                        hidden_states=training_batch_temp.
+                        input_kwargs['hidden_states'],
+                        encoder_hidden_states=training_batch_temp.
+                        input_kwargs['encoder_hidden_states'],
+                        timestep=training_batch_temp.input_kwargs['timestep'],
+                        encoder_hidden_states_image=training_batch_temp.
+                        input_kwargs.get('encoder_hidden_states_image'),
+                        kv_cache=kv_cache1,
+                        crossattn_cache=crossattn_cache,
+                        current_start=current_start_frame * self.frame_seq_length,
+                        start_frame=current_start_frame)
+                    gc.collect()
+                    torch.cuda.empty_cache()
+
                # context_timestep is 0 so we use transformer_2
-                current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer
+                # current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer
                current_model(
                    hidden_states=training_batch_temp.
                    input_kwargs['hidden_states'],
@@ -428,8 +551,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                    timestep=training_batch_temp.input_kwargs['timestep'],
                    encoder_hidden_states_image=training_batch_temp.
                    input_kwargs.get('encoder_hidden_states_image'),
-                    kv_cache=self.kv_cache1,
-                    crossattn_cache=self.crossattn_cache,
+                    kv_cache=kv_cache1,
+                    crossattn_cache=crossattn_cache,
                    current_start=current_start_frame * self.frame_seq_length,
                    start_frame=current_start_frame)

@@ -482,16 +605,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
        else:
            pred_image_or_video_last_21 = pred_image_or_video

-        # Set up gradient mask if we generated more than minimum frames
-        if num_generated_frames != min_num_frames:
-            # Currently, we do not use gradient for the first chunk, since it contains image latents
-            gradient_mask = torch.ones_like(pred_image_or_video_last_21,
-                                            dtype=torch.bool)
-            if self.independent_first_frame:
-                gradient_mask[:, :1] = False
-            else:
-                gradient_mask[:, :self.num_frame_per_block] = False
-
        # Apply gradient masking if needed
        final_output = pred_image_or_video_last_21.to(dtype)
        if gradient_mask is not None:
@@ -520,12 +633,32 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
            training_batch.dmd_latent_vis_dict["min_num_frames"] = torch.tensor(
                min_num_frames, dtype=torch.float32, device=self.device)

-        # Clean up caches
-        assert self.kv_cache1 is not None
-        assert self.crossattn_cache is not None
-        self._reset_simulation_caches(self.kv_cache1, self.crossattn_cache)
+        # Clean up caches - properly free GPU memory
+        self._cleanup_simulation_caches(kv_cache1, crossattn_cache)
+        # Iterate over all buffers in self.transformer and set them to None
+        if self.current_trainstep % self.dfake_gen_update_ratio != 0:
+            for module in self.transformer.modules():
+                for name, buf in module._buffers.items():
+                    if "k_cache" in name or "v_cache" in name:
+                        module._buffers[name] = None
+            if self.transformer_2 is not None:
+                for module in self.transformer_2.modules():
+                    for name, buf in module._buffers.items():
+                        if "k_cache" in name or "v_cache" in name:
+                            module._buffers[name] = None
+        gc.collect()
+        torch.cuda.empty_cache()
+        # assert kv_cache1 is not None
+        # assert crossattn_cache is not None
+        # self._reset_simulation_caches(self.kv_cache1, self.crossattn_cache)

-        return final_output if gradient_mask is not None else pred_image_or_video
+        assert pred_image_or_video.shape[1] == 21, "pred_image_or_video must have 21 frames"
+
+        if gradient_mask is not None:
+            return final_output, exit_timestep
+        else:
+            # return pred_image_or_video, exit_timestep
+            return pred_image_or_video, clean_output, exit_timestep

    def _initialize_simulation_caches(
        self,
@@ -558,6 +691,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
            max_num_frames = num_frames
        num_max_frames = max(max_num_frames, num_frames)
        kv_cache_size = num_max_frames * frame_seq_length
+        assert kv_cache_size == 1560 * 21, "kv_cache_size must be 1560 * 21"

        kv_cache = []
        for _ in range(num_transformer_blocks):
@@ -622,6 +756,30 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                cache_dict["k"].zero_()
                cache_dict["v"].zero_()

+    def _cleanup_simulation_caches(self, kv_cache: list[dict[str, Any]],
+                                  crossattn_cache: list[dict[str, Any]]) -> None:
+        """Properly clean up KV cache and cross-attention cache GPU memory."""
+        if kv_cache is not None:
+            for cache_dict in kv_cache:
+                # Clear tensor references to free GPU memory
+                if "k" in cache_dict and cache_dict["k"] is not None:
+                    cache_dict["k"] = None
+                if "v" in cache_dict and cache_dict["v"] is not None:
+                    cache_dict["v"] = None
+                if "global_end_index" in cache_dict and cache_dict["global_end_index"] is not None:
+                    cache_dict["global_end_index"] = None
+                if "local_end_index" in cache_dict and cache_dict["local_end_index"] is not None:
+                    cache_dict["local_end_index"] = None
+
+        if crossattn_cache is not None:
+            for cache_dict in crossattn_cache:
+                # Clear tensor references to free GPU memory
+                if "k" in cache_dict and cache_dict["k"] is not None:
+                    cache_dict["k"] = None
+                if "v" in cache_dict and cache_dict["v"] is not None:
+                    cache_dict["v"] = None
+                cache_dict["is_init"] = False
+
    def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
        batch = next(self.train_loader_iter, None)  # type: ignore
        if batch is None:
@@ -709,9 +867,40 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                        setattr(batch_gen, key, copy.deepcopy(value))

                generator_loss, gen_log_dict = self.generator_loss(batch_gen)
+                logger.info("train_transformer_2: %s", self.train_transformer_2)
+
                with set_forward_context(current_timestep=batch_gen.timesteps,
                                         attn_metadata=batch_gen.attn_metadata):
                    (generator_loss / gradient_accumulation_steps).backward()
+
+                # Ensure that only one of the two transformers have received gradients
+                if self.train_transformer_2:
+                    # Assert that all gradients are None for transformer
+                    assert all(p.grad is None for p in self.transformer.parameters())
+                    grad_sum = 0
+                    for n, p in self.transformer_2.named_parameters():
+                        if not p.requires_grad:
+                            logger.info("Transformer 2 param %s is not trainable", n)
+                            continue
+                        if p.grad is not None:
+                            grad_sum += p.grad.sum().item()
+                        else:
+                            raise ValueError("Transformer 2 param %s has no gradient", n)
+                    assert grad_sum != 0, "Transformer 2 did not receive gradients"
+                else:
+                    # raise ValueError("Transformer 2 is not training")
+                    assert all(p.grad is None for p in self.transformer_2.parameters())
+                    grad_norm = 0.0
+                    for n, p in self.transformer.named_parameters():
+                        if not p.requires_grad:
+                            logger.info("Transformer param %s is not trainable", n)
+                            continue
+                        if p.grad is not None:
+                            grad_norm += p.grad.data.norm(2).item() ** 2
+                        else:
+                            raise ValueError("Transformer param %s has no gradient", n)
+                    assert grad_norm > 0.0, "Transformer has 0 gradient norm"
+
                total_generator_loss += generator_loss.detach().item()
                generator_log_dict.update(gen_log_dict)
                # Store visualization data from generator training
@@ -723,14 +912,32 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
            if hasattr(
                    self, 'train_transformer_2'
            ) and self.train_transformer_2 and self.transformer_2 is not None:
-                self._clip_model_grad_norm_(batch_gen, self.transformer_2)
+                generator_grad_norm = self._clip_model_grad_norm_(batch_gen, self.transformer_2)
                self.optimizer_2.step()
                self.lr_scheduler_2.step()
            else:
-                self._clip_model_grad_norm_(batch_gen, self.transformer)
+                generator_grad_norm = self._clip_model_grad_norm_(batch_gen, self.transformer)
                self.optimizer.step()
                self.lr_scheduler.step()

+            if (self.current_trainstep >= self.training_args.ema_start_step) and \
+                    (self.generator_ema is None) and (self.training_args.ema_decay > 0):
+                self.generator_ema = EMA_FSDP(
+                    self.transformer, decay=self.training_args.ema_decay)
+                logger.info("Created generator EMA at step %s with decay=%s",
+                            self.current_trainstep, self.training_args.ema_decay)
+
+                # Create EMA for transformer_2 if it exists
+                if self.transformer_2 is not None and self.generator_ema_2 is None:
+                    self.generator_ema_2 = EMA_FSDP(
+                        self.transformer_2, decay=self.training_args.ema_decay)
+                    logger.info(
+                        "Created generator EMA_2 at step %s with decay=%s",
+                        self.current_trainstep, self.training_args.ema_decay)
+                
+                world_group = get_world_group()
+                world_group.barrier()
+
            if self.generator_ema is not None:
                if hasattr(
                        self, 'train_transformer_2'
@@ -748,11 +955,21 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
            world_group.all_reduce(avg_generator_loss,
                                   op=torch.distributed.ReduceOp.AVG)
            training_batch.generator_loss = avg_generator_loss.item()
+            training_batch.generator_grad_norm = generator_grad_norm
+            training_batch.critic_grad_norm = 0.0
+            training_batch.fake_score_loss = 0
+            training_batch.total_loss = training_batch.generator_loss
+            training_batch.fake_score_latent_vis_dict = {
+                "fake_score_timestep": torch.tensor(0.0, device=self.device),
+            }
+            return training_batch
        else:
            training_batch.generator_loss = 0.0
+            training_batch.generator_grad_norm = 0.0

-        logger.debug("Training critic at step %s", self.current_trainstep)
        self.fake_score_optimizer.zero_grad()
+        if self.fake_score_transformer_2 is not None:
+            self.fake_score_optimizer_2.zero_grad()
        total_critic_loss = 0.0
        critic_log_dict = {}

@@ -777,6 +994,27 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
            with set_forward_context(current_timestep=batch_critic.timesteps,
                                     attn_metadata=batch_critic.attn_metadata):
                (critic_loss / gradient_accumulation_steps).backward()
+
+            if self.train_fake_score_transformer_2:
+                assert all(p.grad is None for p in self.fake_score_transformer.parameters())
+                grad_sum = 0
+                for n, p in self.fake_score_transformer_2.named_parameters():
+                    if p.grad is not None:
+                        grad_sum += p.grad.sum().item()
+                    else:
+                        raise ValueError("Fake score transformer 2 param %s has no gradient", n)
+                assert grad_sum != 0, "Fake score transformer 2 did not receive gradients"
+            else:
+                if self.fake_score_transformer_2 is not None:
+                    assert all(p.grad is None for p in self.fake_score_transformer_2.parameters())
+                grad_sum = 0
+                for n, p in self.fake_score_transformer.named_parameters():
+                    if p.grad is not None:
+                        grad_sum += p.grad.sum().item()
+                    else:
+                        raise ValueError("Fake score transformer param %s has no gradient", n)
+                assert grad_sum != 0, "Fake score transformer did not receive gradients"
+
            total_critic_loss += critic_loss.detach().item()
            critic_log_dict.update(crit_log_dict)
            # Store visualization data from critic training
@@ -785,12 +1023,12 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                    batch_critic.fake_score_latent_vis_dict)

        if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None:
-            self._clip_model_grad_norm_(batch_critic,
+            critic_grad_norm = self._clip_model_grad_norm_(batch_critic,
                                        self.fake_score_transformer_2)
            self.fake_score_optimizer_2.step()
            self.fake_score_lr_scheduler_2.step()
        else:
-            self._clip_model_grad_norm_(batch_critic,
+            critic_grad_norm = self._clip_model_grad_norm_(batch_critic,
                                        self.fake_score_transformer)
            self.fake_score_optimizer.step()
            self.fake_score_lr_scheduler.step()
@@ -802,10 +1040,47 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
        world_group.all_reduce(avg_critic_loss,
                               op=torch.distributed.ReduceOp.AVG)
        training_batch.fake_score_loss = avg_critic_loss.item()
-
+        training_batch.critic_grad_norm = critic_grad_norm
        training_batch.total_loss = training_batch.generator_loss + training_batch.fake_score_loss
+
+        gc.collect()
+        torch.cuda.empty_cache()
        return training_batch

+    def _init_negative_prompt(self, training_args: TrainingArgs) -> None:
+        """Initialize negative prompt embeddings for training."""
+        if getattr(self, 'validation_pipeline', None) is None:
+            logger.warning("Validation pipeline not set, skipping negative prompt initialization")
+            return
+
+        logger.info("Initializing negative prompt embeddings...")
+        # Create sampling parameters to get default negative prompt
+        sampling_param = SamplingParam.from_pretrained(training_args.model_path)
+        negative_prompt = sampling_param.negative_prompt
+        
+        batch_negative = ForwardBatch(
+            data_type="video",
+            prompt=negative_prompt,
+            prompt_embeds=[],
+            prompt_attention_mask=[],
+        )
+        
+        # Encode negative prompt
+        # Ensure inference mode is set if needed by pipeline
+        original_inference_mode = getattr(training_args, 'inference_mode', False)
+        training_args.inference_mode = True
+        
+        try:
+            result_batch = self.validation_pipeline.prompt_encoding_stage(
+                batch_negative, training_args)
+            self.negative_prompt_embeds = result_batch.prompt_embeds[0]
+            self.negative_prompt_attention_mask = result_batch.prompt_attention_mask[0]
+            logger.info("Negative prompt embeddings initialized")
+        except Exception as e:
+            logger.error("Failed to initialize negative prompt embeddings: %s", str(e))
+        finally:
+            training_args.inference_mode = original_inference_mode
+
    def _log_training_info(self) -> None:
        """Log self-forcing specific training information."""
        super()._log_training_info()
@@ -815,6 +1090,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
    def visualize_intermediate_latents(self, training_batch: TrainingBatch,
                                       training_args: TrainingArgs, step: int):
        """Add visualization data to wandb logging and save frames to disk."""
+        self.vae = self.vae.to(get_local_torch_device(), dtype=torch.bfloat16)
        wandb_loss_dict = {}

        # Debug logging
@@ -913,6 +1189,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                        video, fps=24, format="mp4")
                    del video, latents

+        self.vae = self.vae.to("cpu")
+
        # Log metadata
        if hasattr(
                training_batch,
@@ -974,8 +1252,16 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
        step_times: deque[float] = deque(maxlen=100)

        self._log_training_info()
-        self._log_validation(self.transformer, self.training_args,
+
+        # Initialize negative prompt regardless of validation settings
+        # self._init_negative_prompt(self.training_args)
+
+        validation_args = deepcopy(self.training_args)
+        self._log_validation(validation_args,
                             self.init_steps)
+        
+        gc.collect()
+        torch.cuda.empty_cache()

        progress_bar = tqdm(
            range(0, self.training_args.max_train_steps),
@@ -1005,20 +1291,23 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
            self.current_trainstep = step
            training_batch.current_vsa_sparsity = current_vsa_sparsity

-            if (step >= self.training_args.ema_start_step) and \
-                    (self.generator_ema is None) and (self.training_args.ema_decay > 0):
-                self.generator_ema = EMA_FSDP(
-                    self.transformer, decay=self.training_args.ema_decay)
-                logger.info("Created generator EMA at step %s with decay=%s",
-                            step, self.training_args.ema_decay)
+            # if (step >= self.training_args.ema_start_step) and \
+            #         (self.generator_ema is None) and (self.training_args.ema_decay > 0):
+            #     self.generator_ema = EMA_FSDP(
+            #         self.transformer, decay=self.training_args.ema_decay)
+            #     logger.info("Created generator EMA at step %s with decay=%s",
+            #                 step, self.training_args.ema_decay)

-                # Create EMA for transformer_2 if it exists
-                if self.transformer_2 is not None and self.generator_ema_2 is None:
-                    self.generator_ema_2 = EMA_FSDP(
-                        self.transformer_2, decay=self.training_args.ema_decay)
-                    logger.info(
-                        "Created generator EMA_2 at step %s with decay=%s",
-                        step, self.training_args.ema_decay)
+            #     # Create EMA for transformer_2 if it exists
+            #     if self.transformer_2 is not None and self.generator_ema_2 is None:
+            #         self.generator_ema_2 = EMA_FSDP(
+            #             self.transformer_2, decay=self.training_args.ema_decay)
+            #         logger.info(
+            #             "Created generator EMA_2 at step %s with decay=%s",
+            #             step, self.training_args.ema_decay)
+                
+            #     world_group = get_world_group()
+            #     world_group.barrier()

            with torch.autocast("cuda", dtype=torch.bfloat16):
                training_batch = self.train_one_step(training_batch)
@@ -1026,7 +1315,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
            total_loss = training_batch.total_loss
            generator_loss = training_batch.generator_loss
            fake_score_loss = training_batch.fake_score_loss
-            grad_norm = training_batch.grad_norm
+            generator_grad_norm = training_batch.generator_grad_norm
+            critic_grad_norm = training_batch.critic_grad_norm

            step_time = time.perf_counter() - start_time
            step_times.append(step_time)
@@ -1041,8 +1331,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                f"{fake_score_loss:.4f}",
                "step_time":
                f"{step_time:.2f}s",
-                "grad_norm":
-                grad_norm,
+                "generator_grad_norm":
+                generator_grad_norm,
+                "critic_grad_norm":
+                critic_grad_norm,
                "ema":
                "✓" if (self.generator_ema is not None and self.is_ema_ready())
                else "✗",
@@ -1066,8 +1358,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                    step_time,
                    "avg_step_time":
                    avg_step_time,
-                    "grad_norm":
-                    grad_norm,
+                    "generator_grad_norm":
+                    generator_grad_norm,
+                    "critic_grad_norm":
+                    critic_grad_norm,
                }
                if (step % self.dfake_gen_update_ratio == 0):
                    log_data["train_generator_loss"] = generator_loss
@@ -1147,6 +1441,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):

                if self.transformer:
                    self.transformer.train()
+                if getattr(self, 'transformer_2', None) is not None:
+                    self.transformer_2.train()
                self.sp_group.barrier()

            if (self.training_args.weight_only_checkpointing_steps > 0
@@ -1179,11 +1475,15 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                                                   None),
                    generator_ema_2=getattr(self, 'generator_ema_2', None))

-                if self.training_args.use_ema and self.is_ema_ready():
-                    self.save_ema_weights(self.training_args.output_dir, step)
+                # if self.training_args.use_ema and self.is_ema_ready():
+                #     self.save_ema_weights(self.training_args.output_dir, step)
+
+                self.sp_group.barrier()

            if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
-                self._log_validation(self.transformer, self.training_args, step)
+                self._log_validation(deepcopy(self.training_args), step)
+                gc.collect()
+                torch.cuda.empty_cache()

        wandb.finish()

@@ -1217,9 +1517,10 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
                                           None),
            generator_ema_2=getattr(self, 'generator_ema_2', None))

-        if self.training_args.use_ema and self.is_ema_ready():
-            self.save_ema_weights(self.training_args.output_dir,
-                                  self.training_args.max_train_steps)
+        # if self.training_args.use_ema and self.is_ema_ready():
+        #     self.save_ema_weights(self.training_args.output_dir,
+        #                           self.training_args.max_train_steps)

        if get_sp_group():
            cleanup_dist_env_and_memory()
+
@@ -102,24 +102,20 @@ class TrainingPipeline(LoRAPipeline, ABC):
        self.transformer_2 = self.get_module("transformer_2", None)
        self.seed = training_args.seed
        self.set_schemas()
+        self.set_trainable()

        # Set random seeds for deterministic training
        assert self.seed is not None, "seed must be set"
        set_random_seed(self.seed)
        self.transformer.train()
+        self.transformer.requires_grad_(True)
        if training_args.enable_gradient_checkpointing_type is not None:
            self.transformer = apply_activation_checkpointing(
                self.transformer,
                checkpointing_type=training_args.
                enable_gradient_checkpointing_type)
-            if self.transformer_2 is not None:
-                self.transformer_2 = apply_activation_checkpointing(
-                    self.transformer_2,
-                    checkpointing_type=training_args.
-                    enable_gradient_checkpointing_type)

        noise_scheduler = self.modules["scheduler"]
-        self.set_trainable()
        params_to_optimize = self.transformer.parameters()
        params_to_optimize = list(
            filter(lambda p: p.requires_grad, params_to_optimize))
@@ -152,6 +148,11 @@ class TrainingPipeline(LoRAPipeline, ABC):
            # Ensure transformer_2 has trainable parameters before creating optimizer
            self.transformer_2.train()
            self.transformer_2.requires_grad_(True)
+            if training_args.enable_gradient_checkpointing_type is not None:
+                self.transformer_2 = apply_activation_checkpointing(
+                    self.transformer_2,
+                    checkpointing_type=training_args.
+                    enable_gradient_checkpointing_type)
            params_to_optimize_2 = self.transformer_2.parameters()
            params_to_optimize_2 = list(
                filter(lambda p: p.requires_grad, params_to_optimize_2))
@@ -232,17 +233,20 @@ class TrainingPipeline(LoRAPipeline, ABC):
    def _enable_training(self, model: torch.nn.Module,
                         optimizer: torch.optim.Optimizer) -> None:
        """Enable training mode and gradients for the specified model."""
-        for param in model.parameters():
-            param.requires_grad = True
-        model.train()
-        optimizer.zero_grad()
+        # for param in model.parameters():
+        #     param.requires_grad = True
+        # model.train()
+        # optimizer.zero_grad()
+        return

    def _disable_training(self, model: torch.nn.Module,
                          optimizer: torch.optim.Optimizer) -> None:
        """Disable training mode and gradients for the specified model."""
-        for param in model.parameters():
-            param.requires_grad = False
-        optimizer.zero_grad(set_to_none=True)
+        # for param in model.parameters():
+        #     param.requires_grad = False
+        # model.eval()
+        # optimizer.zero_grad(set_to_none=True)
+        return

    def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
        batch = next(self.train_loader_iter, None)  # type: ignore
@@ -357,6 +361,20 @@ class TrainingPipeline(LoRAPipeline, ABC):
        #     pass

        indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
+        # timestep = self.noise_scheduler.timesteps[indices].to(device=device)
+
+        # if timestep < self.training_args.boundary_ratio * self.noise_scheduler.config.num_train_timesteps:
+        #     self.train_transformer_2 = True
+        # else:
+        #     self.train_transformer_2 = False
+
+        # # Broadcast the decision to all processes
+        # decision = torch.tensor(1.0 if self.train_transformer_2 else 0.0,
+        #                         device=self.device)
+        # dist.broadcast(decision, src=0)
+        # dist.broadcast(timestep, src=0)
+        # self.train_transformer_2 = decision.item() == 1.0
+
        return self.noise_scheduler.timesteps[indices].to(device=device)

    def _build_attention_metadata(
@@ -698,6 +716,8 @@ class TrainingPipeline(LoRAPipeline, ABC):
                                  validation_batch: dict[str, Any],
                                  num_inference_steps: int) -> ForwardBatch:
        sampling_param.prompt = validation_batch['prompt']
+        if 'image_path' in validation_batch:
+            sampling_param.image_path = validation_batch['image_path']
        sampling_param.height = training_args.num_height
        sampling_param.width = training_args.num_width
        sampling_param.num_inference_steps = num_inference_steps
@@ -37,7 +37,7 @@ def gather_state_dict_on_cpu_rank0(
    sharded_sd = model.state_dict()
    param_requires_grad = set([
        k.replace("._checkpoint_wrapped_module.", ".")
-        for k, v in dict(model.named_parameters()).items() if v.requires_grad
+        for k, v in dict(model.named_parameters()).items()
    ])
    for param_name, param in sharded_sd.items():
        if param_name not in param_requires_grad:
@@ -255,8 +255,6 @@ def save_distillation_checkpoint(
        if generator_scheduler is not None:
            generator_states["scheduler"] = SchedulerWrapper(
                generator_scheduler)
-        if generator_ema is not None:
-            generator_states["ema"] = generator_ema.state_dict()

        generator_dcp_dir = os.path.join(save_dir, "distributed_checkpoint",
                                         "generator")
@@ -288,8 +286,6 @@ def save_distillation_checkpoint(
            if generator_scheduler_2 is not None:
                generator_2_states["scheduler"] = SchedulerWrapper(
                    generator_scheduler_2)
-            if generator_ema_2 is not None:
-                generator_2_states["ema"] = generator_ema_2.state_dict()

            generator_2_dcp_dir = os.path.join(save_dir,
                                               "distributed_checkpoint",
@@ -415,6 +411,69 @@ def save_distillation_checkpoint(
            rank,
            local_main_process_only=False)

+    # Persist EMA separately to avoid shape mismatches across ranks.
+    # Supports:
+    #   - mode="rank0_full": save consolidated EMA only on rank 0
+    #   - mode="local_shard": save per-rank EMA shard for each rank
+    try:
+        if generator_ema is not None and getattr(generator_ema, "mode",
+                                                 None) == "rank0_full":
+            _save_rank0_full_ema_safetensors(generator_ema,
+                                             generator_transformer, rank,
+                                             save_dir, "generator_ema")
+        elif generator_ema is not None and getattr(generator_ema, "mode",
+                                                   None) == "local_shard":
+            # Save per-rank shard
+            ema_dir_shard = os.path.join(save_dir, "ema_local_shard")
+            os.makedirs(ema_dir_shard, exist_ok=True)
+            ema_shard_path = os.path.join(ema_dir_shard, f"generator_ema_rank{rank}.pt")
+            torch.save(generator_ema.state_dict(), ema_shard_path)
+            logger.info(
+                "rank: %s, saved generator EMA shard (local_shard) to %s",
+                rank,
+                ema_shard_path,
+                local_main_process_only=False)
+
+            # Also consolidate EMA to a single full-state file on rank 0 by applying EMA to the model and gathering
+            # _consolidate_local_shard_ema_and_save_safetensors(
+            #     generator_ema, generator_transformer, rank, save_dir,
+            #     "generator_ema")
+
+    except Exception as e:
+        logger.warning("rank: %s, failed saving EMA separately: %s",
+                       rank,
+                       str(e))
+
+    try:
+        if generator_ema_2 is not None and getattr(generator_ema_2, "mode",
+                                                   None) == "rank0_full":
+            _save_rank0_full_ema_safetensors(generator_ema_2,
+                                             generator_transformer_2, rank,
+                                             save_dir, "generator_ema_2")
+        elif generator_ema_2 is not None and getattr(generator_ema_2, "mode",
+                                                     None) == "local_shard":
+            # Save per-rank shard for EMA_2
+            ema_dir_shard_2 = os.path.join(save_dir, "ema_local_shard")
+            os.makedirs(ema_dir_shard_2, exist_ok=True)
+            ema2_shard_path = os.path.join(ema_dir_shard_2, f"generator_ema_2_rank{rank}.pt")
+            torch.save(generator_ema_2.state_dict(), ema2_shard_path)
+            logger.info(
+                "rank: %s, saved generator_2 EMA shard (local_shard) to %s",
+                rank,
+                ema2_shard_path,
+                local_main_process_only=False)
+
+            # Also consolidate EMA_2 to a single full-state file on rank 0
+            # _consolidate_local_shard_ema_and_save_safetensors(
+            #     generator_ema_2, generator_transformer_2, rank, save_dir,
+            #     "generator_ema_2")
+    except Exception as e:
+        logger.warning("rank: %s, failed saving EMA_2 separately: %s",
+                       rank,
+                       str(e))
+
+    return
+
    # Save generator model weights (consolidated) for inference
    cpu_state = gather_state_dict_on_cpu_rank0(generator_transformer,
                                               device=None)
@@ -452,46 +511,46 @@ def save_distillation_checkpoint(
        logger.info("--> distillation checkpoint saved at step %s to %s", step,
                    weight_path)

-        # Save generator_2 model weights (consolidated) for inference (MoE support)
-        if generator_transformer_2 is not None:
-            inference_save_dir_2 = os.path.join(
-                save_dir, "generator_2_inference_transformer")
-            cpu_state_2 = gather_state_dict_on_cpu_rank0(
-                generator_transformer_2, device=None)
+    # Save generator_2 model weights (consolidated) for inference (MoE support)
+    if generator_transformer_2 is not None:
+        inference_save_dir_2 = os.path.join(
+            save_dir, "generator_2_inference_transformer")
+        cpu_state_2 = gather_state_dict_on_cpu_rank0(
+            generator_transformer_2, device=None)

-            if rank == 0:
-                os.makedirs(inference_save_dir_2, exist_ok=True)
-                weight_path_2 = os.path.join(
-                    inference_save_dir_2, "diffusion_pytorch_model.safetensors")
-                logger.info(
-                    "rank: %s, saving consolidated generator_2 inference checkpoint to %s",
-                    rank,
-                    weight_path_2,
-                    local_main_process_only=False)
+        if rank == 0:
+            os.makedirs(inference_save_dir_2, exist_ok=True)
+            weight_path_2 = os.path.join(
+                inference_save_dir_2, "diffusion_pytorch_model.safetensors")
+            logger.info(
+                "rank: %s, saving consolidated generator_2 inference checkpoint to %s",
+                rank,
+                weight_path_2,
+                local_main_process_only=False)

-                # Convert training format to diffusers format and save
-                diffusers_state_dict_2 = custom_to_hf_state_dict(
-                    cpu_state_2,
-                    generator_transformer_2.reverse_param_names_mapping)
-                save_file(diffusers_state_dict_2, weight_path_2)
+            # Convert training format to diffusers format and save
+            diffusers_state_dict_2 = custom_to_hf_state_dict(
+                cpu_state_2,
+                generator_transformer_2.reverse_param_names_mapping)
+            save_file(diffusers_state_dict_2, weight_path_2)

-                logger.info(
-                    "rank: %s, consolidated generator_2 inference checkpoint saved to %s",
-                    rank,
-                    weight_path_2,
-                    local_main_process_only=False)
+            logger.info(
+                "rank: %s, consolidated generator_2 inference checkpoint saved to %s",
+                rank,
+                weight_path_2,
+                local_main_process_only=False)

-                # Save model config
-                config_dict_2 = generator_transformer_2.hf_config
-                if "dtype" in config_dict_2:
-                    del config_dict_2["dtype"]  # TODO
-                config_path_2 = os.path.join(inference_save_dir_2,
-                                             "config.json")
-                with open(config_path_2, "w") as f:
-                    json.dump(config_dict_2, f, indent=4)
-                logger.info(
-                    "--> generator_2 distillation checkpoint saved at step %s to %s",
-                    step, weight_path_2)
+            # Save model config
+            config_dict_2 = generator_transformer_2.hf_config
+            if "dtype" in config_dict_2:
+                del config_dict_2["dtype"]  # TODO
+            config_path_2 = os.path.join(inference_save_dir_2,
+                                            "config.json")
+            with open(config_path_2, "w") as f:
+                json.dump(config_dict_2, f, indent=4)
+            logger.info(
+                "--> generator_2 distillation checkpoint saved at step %s to %s",
+                step, weight_path_2)


 def load_checkpoint(transformer,
@@ -641,18 +700,38 @@ def load_distillation_checkpoint(
        end_time - begin_time,
        local_main_process_only=False)

-    # Load EMA state if available and generator_ema is provided
+    # Load EMA separately if saved in rank0_full mode
    if generator_ema is not None:
        try:
-            ema_state = generator_states.get("ema")
-            if ema_state is not None:
-                generator_ema.load_state_dict(ema_state)
-                logger.info("rank: %s, generator EMA state loaded successfully",
-                            rank)
-            else:
-                logger.info("rank: %s, no EMA state found in checkpoint", rank)
+            if getattr(generator_ema, "mode", None) == "rank0_full":
+                ema_path = os.path.join(checkpoint_path, "ema",
+                                        "generator_ema.pt")
+                if rank == 0 and os.path.exists(ema_path):
+                    ema_state = torch.load(ema_path, map_location="cpu")
+                    generator_ema.load_state_dict(ema_state)
+                    logger.info(
+                        "rank: %s, generator EMA (rank0_full) loaded from %s",
+                        rank, ema_path)
+                elif rank == 0:
+                    logger.info(
+                        "rank: %s, generator EMA file not found at %s; skipping",
+                        rank, ema_path)
+            elif getattr(generator_ema, "mode", None) == "local_shard":
+                ema_path = os.path.join(checkpoint_path, "ema_local_shard",
+                                        f"generator_ema_rank{rank}.pt")
+                if os.path.exists(ema_path):
+                    ema_state = torch.load(ema_path, map_location="cpu")
+                    generator_ema.load_state_dict(ema_state)
+                    logger.info(
+                        "rank: %s, generator EMA shard (local_shard) loaded from %s",
+                        rank, ema_path)
+                else:
+                    logger.info(
+                        "rank: %s, generator EMA shard file not found at %s; skipping",
+                        rank, ema_path)
+
        except Exception as e:
-            logger.warning("rank: %s, failed to load EMA state: %s", rank,
+            logger.warning("rank: %s, failed to load generator EMA: %s", rank,
                           str(e))

    # Load generator_2 distributed checkpoint (MoE support)
@@ -692,26 +771,45 @@ def load_distillation_checkpoint(
                end_time - begin_time,
                local_main_process_only=False)

-            # Load EMA_2 state if available and generator_ema_2 is provided
+            # Load EMA separately if saved in rank0_full mode
            if generator_ema_2 is not None:
                try:
-                    ema_2_state = generator_2_states.get("ema")
-                    if ema_2_state is not None:
-                        generator_ema_2.load_state_dict(ema_2_state)
-                        logger.info(
-                            "rank: %s, generator_2 EMA state loaded successfully",
-                            rank)
-                    else:
-                        logger.info(
-                            "rank: %s, no EMA_2 state found in checkpoint",
-                            rank)
+                    if getattr(generator_ema_2, "mode", None) == "rank0_full":
+                        ema_path = os.path.join(checkpoint_path, "ema",
+                                                "generator_ema_2.pt")
+                        if rank == 0 and os.path.exists(ema_path):
+                            ema_state = torch.load(ema_path, map_location="cpu")
+                            generator_ema.load_state_dict(ema_state)
+                            logger.info(
+                                "rank: %s, generator EMA 2 (rank0_full) loaded from %s",
+                                rank, ema_path)
+                        elif rank == 0:
+                            logger.info(
+                                "rank: %s, generator EMA 2 file not found at %s; skipping",
+                                rank, ema_path)
+                    elif getattr(generator_ema_2, "mode", None) == "local_shard":
+                        ema_path = os.path.join(checkpoint_path, "ema_local_shard",
+                                                f"generator_ema_2_rank{rank}.pt")
+                        if os.path.exists(ema_path):
+                            ema_state = torch.load(ema_path, map_location="cpu")
+                            generator_ema_2.load_state_dict(ema_state)
+                            logger.info(
+                                "rank: %s, generator EMA 2 shard (local_shard) loaded from %s",
+                                rank, ema_path)
+                        else:
+                            logger.info(
+                                "rank: %s, generator EMA 2 shard file not found at %s; skipping",
+                                rank, ema_path)
+
                except Exception as e:
-                    logger.warning("rank: %s, failed to load EMA_2 state: %s",
-                                   rank, str(e))
+                    logger.warning("rank: %s, failed to load generator EMA 2: %s", rank,
+                                str(e))
        else:
            logger.info("rank: %s, generator_2 checkpoint not found, skipping",
                        rank)

+    # return step
+
    # Load critic distributed checkpoint
    critic_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint",
                                  "critic")
@@ -1153,14 +1251,104 @@ def custom_to_hf_state_dict(
    return new_state_dict


+def _save_full_ema_safetensors_from_state(
+    state_dict: dict[str, Any],
+    reverse_param_names_mapping: dict[str, tuple[str, int, int]],
+    output_path: str,
+) -> None:
+    """
+    Convert a training-format state_dict to HF format and save as safetensors.
+    """
+    diffusers_state_dict = custom_to_hf_state_dict(state_dict,
+                                                   reverse_param_names_mapping)
+    save_file(diffusers_state_dict, output_path)
+
+
+def _save_rank0_full_ema_safetensors(
+    ema: "EMA_FSDP",
+    module,
+    rank: int,
+    save_dir: str,
+    base_name: str,
+) -> None:
+    if rank != 0:
+        return
+    ema_dir = os.path.join(save_dir, "ema")
+    os.makedirs(ema_dir, exist_ok=True)
+    output_path = os.path.join(ema_dir, f"{base_name}.safetensors")
+    ema_state = ema.state_dict()
+    _save_full_ema_safetensors_from_state(ema_state,
+                                          module.reverse_param_names_mapping,
+                                          output_path)
+    logger.info(
+        "rank: %s, saved %s as consolidated EMA safetensors to %s",
+        rank,
+        base_name,
+        output_path,
+        local_main_process_only=False)
+
+
+def _consolidate_local_shard_ema_and_save_safetensors(
+    ema: "EMA_FSDP",
+    module,
+    rank: int,
+    save_dir: str,
+    base_name: str,
+) -> None:
+    try:
+        # Temporarily apply EMA to the live (sharded) module and gather full CPU state on rank 0
+        with ema.apply_to_model(module):
+            cpu_state_full = gather_state_dict_on_cpu_rank0(module, device=None)
+        if rank == 0:
+            ema_dir = os.path.join(save_dir, "ema")
+            os.makedirs(ema_dir, exist_ok=True)
+            output_path = os.path.join(ema_dir, f"{base_name}.safetensors")
+            _save_full_ema_safetensors_from_state(
+                cpu_state_full, module.reverse_param_names_mapping, output_path)
+            logger.info(
+                "rank: %s, saved consolidated %s EMA (from local_shard) as safetensors to %s",
+                rank,
+                base_name,
+                output_path,
+                local_main_process_only=False)
+    except Exception as ce:
+        logger.warning(
+            "rank: %s, failed consolidating %s EMA (local_shard): %s",
+            rank,
+            base_name,
+            str(ce))
+
 def shift_timestep(timestep: torch.Tensor, shift: float,
-                   num_train_timestep: float) -> torch.Tensor:
+                   min_timestep: int = 0, max_timestep: int = 1000) -> torch.Tensor:
    if shift == 1:
        return timestep
-    t = timestep / num_train_timestep
+    t = (timestep - min_timestep) / (max_timestep - min_timestep)
    denominator = 1 + (shift - 1) * t
-    return num_train_timestep * (shift * t / denominator)
+    return min_timestep + (max_timestep - min_timestep) * (shift * t / denominator)

+def unshift_timestep(shifted_timestep: torch.Tensor, shift: float,
+                     min_timestep: int = 0, max_timestep: int = 1000) -> torch.Tensor:
+    """Reverse the shift_timestep operation.
+    
+    Given a timestep that has been shifted by shift_timestep, this function
+    recovers the original timestep.
+    
+    Args:
+        shifted_timestep: The shifted timestep tensor
+        shift: The shift value that was used in shift_timestep
+        min_timestep: Minimum timestep value (default: 0)
+        max_timestep: Maximum timestep value (default: 1000)
+    
+    Returns:
+        The original (unshifted) timestep tensor
+    """
+    if shift == 1:
+        return shifted_timestep
+    s = (shifted_timestep - min_timestep) / (max_timestep - min_timestep)
+    # Solve for original normalized timestep t: s = shift * t / (1 + (shift - 1) * t)
+    # Rearranging: t = s / (shift - s * (shift - 1))
+    t = s / (shift - s * (shift - 1))
+    return min_timestep + (max_timestep - min_timestep) * t

 # coding=utf-8
 # Copyright 2025 The HuggingFace Inc. team.
@@ -1616,7 +1804,8 @@ def _local_numel(p: torch.Tensor) -> int:
 def count_trainable(model: torch.nn.Module) -> int:
    return sum(_local_numel(p) for p in model.parameters() if p.requires_grad)

-
+import gc
+from torch.distributed.tensor import DTensor
 class EMA_FSDP:
    """
    FSDP2-friendly EMA with two modes:
@@ -1676,8 +1865,11 @@ class EMA_FSDP:
        self.shadow = {}
        for name, p in module.named_parameters():
            if not p.requires_grad:
+                logger.info(f"Parameter {name} is not trainable in ema init")
                continue
            local = self._to_local_tensor(p.detach())
+            # if "scale_shift_table" in name or "proj_out" in name:
+            #     logger.info(f"{name} in ema init is DTensor? {isinstance(p, DTensor)} with parameter shape {p.shape} and local shape {local.shape}")
            self.shadow[name] = local.clone().float().cpu()

    @torch.no_grad()
@@ -1698,13 +1890,24 @@ class EMA_FSDP:
        # local_shard: update local shard EMA on every rank
        for name, p in module.named_parameters():
            if not p.requires_grad:
+                logger.info(f"Parameter {name} is not trainable in ema update")
                continue
            local = self._to_local_tensor(p.detach())
            v_cpu = local.float().cpu()
+            # if "scale_shift_table" in name or "proj_out" in name:
+            #     logger.info(f"{name} in ema update is DTensor? {isinstance(p, DTensor)} with parameter shape {p.shape} and local shape {local.shape}")
            if name not in self.shadow:
+                logger.info(f"Parameter {name} is not in shadow in ema update")
                self.shadow[name] = v_cpu.clone()
            else:
-                self.shadow[name].mul_(d).add_(v_cpu, alpha=1.0 - d)
+                if self.shadow[name].shape != v_cpu.shape:
+                    logger.info(f"Parameter {name} shape mismatch in ema update. EMA shape {self.shadow[name].shape}, local shape {v_cpu.shape}")
+                    self.shadow[name] = v_cpu.clone()
+                else:
+                    self.shadow[name].mul_(d).add_(v_cpu, alpha=1.0 - d)
+            del local
+        gc.collect()
+        torch.cuda.empty_cache()

    def state_dict(self) -> dict[str, torch.Tensor]:
        if self.mode == "rank0_full":
@@ -1746,6 +1949,7 @@ class EMA_FSDP:
            with torch.no_grad():
                for name, p in self.module.named_parameters():
                    if not p.requires_grad:
+                        logger.info(f"Parameter {name} is not trainable in ema apply_to_model")
                        continue
                    # Save local shard
                    p_local = EMA_FSDP._to_local_tensor(p.detach())
@@ -1757,6 +1961,7 @@ class EMA_FSDP:
                    if name in self.ema.shadow:
                        ema_cpu = self.ema.shadow[name]
                        if ema_cpu.numel() != p_local.numel():
+                            logger.info(f"Parameter {name} shape mismatch in ema apply_to_model. EMA shape {ema_cpu.shape}, local shape {p_local.shape}")
                            # Shard shape mismatch (e.g., empty shard here), skip
                            continue
                        # Copy EMA shard into local param shard
@@ -1774,10 +1979,11 @@ class EMA_FSDP:
                            continue
                        saved_local = self.saved[name]
                        if saved_local.numel() != p_local.numel():
+                            logger.info(f"Parameter {name} shape mismatch in ema apply_to_model exit. EMA shape {saved_local.shape}, local shape {p_local.shape}")
                            continue
                        p_local.copy_(saved_local)
            self.saved.clear()
            return False

    def apply_to_model(self, module):
-        return EMA_FSDP._ApplyEMACtx(self, module)
+        return EMA_FSDP._ApplyEMACtx(self, module)
@@ -42,7 +42,7 @@ class WanDistillationPipeline(DistillationPipeline):
            training_args.model_path,
            args=args_copy,  # type: ignore
            inference_mode=True,
-            loaded_modules={"transformer": self.get_module("transformer")},
+            loaded_modules={"transformer": self.get_module("transformer"), "transformer_2": self.get_module("transformer_2")},
            tp_size=training_args.tp_size,
            sp_size=training_args.sp_size,
            num_gpus=training_args.num_gpus,
@@ -49,7 +49,7 @@ class WanSelfForcingDistillationPipeline(SelfForcingDistillationPipeline):
            sp_size=training_args.sp_size,
            num_gpus=training_args.num_gpus,
            pin_cpu_memory=training_args.pin_cpu_memory,
-            dit_cpu_offload=True)
+            dit_cpu_offload=False)

        self.validation_pipeline = validation_pipeline

@@ -42,6 +42,7 @@ class WanTrainingPipeline(TrainingPipeline):
            inference_mode=True,
            loaded_modules={
                "transformer": self.get_module("transformer"),
+                "transformer_2": self.get_module("transformer_2"),
            },
            tp_size=training_args.tp_size,
            sp_size=training_args.sp_size,
Author	SHA1	Message	Date
JerryZhou54	96f3cbd487	add 1333 script	2026-03-02 21:50:31 +00:00
JerryZhou54	4276b0567f	ckpt	2026-02-09 22:41:25 +00:00
JerryZhou54	36515d5b72	Resolve OOM	2025-11-27 21:08:32 +00:00
JerryZhou54	e02f9d8d47	setup training for 14b MoE	2025-11-25 09:33:33 +00:00
JerryZhou54	df7381e2b0	Support sampling from the full timestep range in dmd_forward & faker_score_forward when the output video is only generated by high noise generator	2025-10-13 21:10:34 +00:00
JerryZhou54	de11ea3020	Small runtime error	2025-10-11 09:58:20 +00:00
JerryZhou54	4a527bccf2	Add new add_noise_high logic to support full MoE training	2025-10-11 09:21:08 +00:00
JerryZhou54	dabf87fefa	Add config for 1.3B MoE	2025-10-11 04:18:03 +00:00
JerryZhou54	8996a391a1	Add pred_noise_to_x_bound & add_noise_high	2025-10-11 00:29:31 +00:00
JerryZhou54	fba8b61c4f	Fix MoE recipe and Add 1.3B MoE script	2025-10-08 03:58:07 +00:00
JerryZhou54	55d8c1e5fb	Fix small runtime issues	2025-10-07 05:03:36 +00:00
JerryZhou54	aeb8f2e5ac	Add matthew's timestep change & add real score guidance scale 2	2025-10-06 22:24:44 +00:00
JerryZhou54	f755dd1ad5	Refactor sf distill code	2025-10-05 23:46:35 +00:00
JerryZhou54	256f788d34	checkpoint	2025-10-04 23:56:36 +00:00