Update DF.py

DF.py

@@ -18,7 +18,7 @@ def generate_video(prompt, negative_prompt, width, height, num_frames,
     start_time = time.time()
     torch.cuda.empty_cache()
 
-    # Load model
+    # Load VAE and Wan2.2 pipeline
     vae = AutoencoderKLWan.from_pretrained(
         "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
         subfolder="vae",
@@ -31,15 +31,9 @@ def generate_video(prompt, negative_prompt, width, height, num_frames,
         torch_dtype=torch.bfloat16,
     )
 
-    # Load DFloat11 optimizations
+    # Only apply second-stage DFloat11 model
     DFloat11Model.from_pretrained(
-        "LeanModels/Wan2.2-T2V-A14B-DF11",
-        device="cpu",
-        cpu_offload=cpu_offload,
-        bfloat16_model=pipe.transformer,
-    )
-    DFloat11Model.from_pretrained(
-        "LeanModels/Wan2.2-T2V-A14B-2-DF11",
+        "DFloat11/Wan2.2-T2V-A14B-2-DF11",
         device="cpu",
         cpu_offload=cpu_offload,
         bfloat16_model=pipe.transformer_2,
@@ -70,7 +64,7 @@ def generate_video(prompt, negative_prompt, width, height, num_frames,
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🎬 Wan2.2 + DFloat11 - Text to Video Generator")
+    gr.Markdown("## 🎬 Wan2.2 + DFloat11 (Stage 2 only) - Text to Video Generator")
 
     with gr.Row():
         prompt = gr.Textbox(label="Prompt", value="A serene koi pond at night, with glowing lanterns reflecting on the rippling water. Ethereal fireflies dance above as cherry blossoms gently fall.", lines=3)