Update optimization.py

optimization.py

@@ -42,140 +42,53 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
 
     @spaces.GPU(duration=1500)
     def compile_transformer():
-        
-        # pipeline.load_lora_weights(
-        #     "Kijai/WanVideo_comfy", 
-        #     weight_name="Lightx2v/lightx2v_T2V_14B_cfg_step_distill_v2_lora_rank128_bf16.safetensors", 
-        #     adapter_name="lightning"
-        # )
-        # kwargs_lora = {}
-        # kwargs_lora["load_into_transformer_2"] = True
-        # pipeline.load_lora_weights(
-        #     "Kijai/WanVideo_comfy", 
-        #      weight_name="Lightx2v/lightx2v_T2V_14B_cfg_step_distill_v2_lora_rank128_bf16.safetensors", 
-        #     #weight_name="Wan22-Lightning/Wan2.2-Lightning_T2V-A14B-4steps-lora_LOW_fp16.safetensors", 
-        #     adapter_name="lightning_2", **kwargs_lora
-        # )
-        # pipeline.set_adapters(["lightning", "lightning_2"], adapter_weights=[1., 1.])
-
-
-        # pipeline.fuse_lora(adapter_names=["lightning"], lora_scale=3., components=["transformer"])
-        # pipeline.fuse_lora(adapter_names=["lightning_2"], lora_scale=1., components=["transformer_2"])
-        # pipeline.unload_lora_weights()
-
-
-
-        pipeline.load_lora_weights(
-            "Kijai/WanVideo_comfy", 
-            weight_name="Lightx2v/lightx2v_T2V_14B_cfg_step_distill_v2_lora_rank128_bf16.safetensors", 
-            adapter_name="lightning"
-        )
-        kwargs_lora = {}
-        kwargs_lora["load_into_transformer_2"] = True
-        # pipeline.load_lora_weights(
-        #     #"drozbay/Wan2.2_A14B_lora_extract", 
-        #      "Kijai/WanVideo_comfy", 
-        #     #weight_name="MTVCrafter/Wan2_1_MTV-Crafter_motion_adapter_bf16.safetensors", 
-        #      #weight_name="Skyreels/Wan2_1_Skyreels-v2-T2V-720P_LoRA_rank_64_fp16.safetensors",
-        #     #weight_name="Pusa/Wan21_PusaV1_LoRA_14B_rank512_bf16.safetensors", 
-        #     weight_name="Wan22-Lightning/Wan2.2-Lightning_T2V-A14B-4steps-lora_LOW_fp16.safetensors", 
-        #     adapter_name="lightning_2", **kwargs_lora
-        # )
-
-
-
-        pipeline.load_lora_weights(
-            #"drozbay/Wan2.2_A14B_lora_extract", 
-             "lightx2v/Wan2.2-Lightning", 
-            #weight_name="MTVCrafter/Wan2_1_MTV-Crafter_motion_adapter_bf16.safetensors", 
-             #weight_name="Skyreels/Wan2_1_Skyreels-v2-T2V-720P_LoRA_rank_64_fp16.safetensors",
-            weight_name="Wan2.2-T2V-A14B-4steps-lora-rank64-Seko-V1/low_noise_model.safetensors", 
-            #weight_name="Wan2.2-T2V-A14B-4steps-lora-rank64-Seko-V1.1/low_noise_model.safetensors", 
-            adapter_name="lightning_2", **kwargs_lora
-        )
-
-        # pipeline.load_lora_weights(
-        #     #"drozbay/Wan2.2_A14B_lora_extract", 
-        #      "ostris/wan22_i2v_14b_orbit_shot_lora", 
-        #     #weight_name="MTVCrafter/Wan2_1_MTV-Crafter_motion_adapter_bf16.safetensors", 
-        #      #weight_name="Skyreels/Wan2_1_Skyreels-v2-T2V-720P_LoRA_rank_64_fp16.safetensors",
-        #     weight_name="wan22_14b_i2v_orbit_low_noise.safetensors", 
-        #     #weight_name="Wan2.2-T2V-A14B-4steps-lora-rank64-Seko-V1.1/low_noise_model.safetensors", 
-        #     adapter_name="lightning_200", **kwargs_lora
-        # )
-
-
-    
-        pipeline.load_lora_weights(
-            #"drozbay/Wan2.2_A14B_lora_extract", 
-             "deadman44/Wan2.2_T2i_T2v_LoRA", 
-            #weight_name="MTVCrafter/Wan2_1_MTV-Crafter_motion_adapter_bf16.safetensors", 
-             #weight_name="Skyreels/Wan2_1_Skyreels-v2-T2V-720P_LoRA_rank_64_fp16.safetensors",
-            weight_name="lora_wan2.2_myjd_Low_v01.safetensors", 
-            #weight_name="Wan2.2-T2V-A14B-4steps-lora-rank64-Seko-V1.1/low_noise_model.safetensors", 
-            adapter_name="lightning_22", **kwargs_lora
-        )
-
-
-
-
-
-
-
-        pipeline.set_adapters(["lightning", "lightning_2","lightning_22"], adapter_weights=[1., 1.,1.])
-
-
-        pipeline.fuse_lora(adapter_names=["lightning"], lora_scale=3., components=["transformer"])
-        pipeline.fuse_lora(adapter_names=["lightning_2"], lora_scale=1., components=["transformer_2"])
-        pipeline.fuse_lora(adapter_names=["lightning_22"], lora_scale=1., components=["transformer_2"])
-       
-        pipeline.unload_lora_weights()
-        
-        with capture_component_call(pipeline, 'transformer') as call:
-            pipeline(*args, **kwargs)
-        
-        dynamic_shapes = tree_map_only((torch.Tensor, bool), lambda t: None, call.kwargs)
-        dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
-
-        quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
-        quantize_(pipeline.transformer_2, Float8DynamicActivationFloat8WeightConfig())
-        
-        hidden_states: torch.Tensor = call.kwargs['hidden_states']
-        hidden_states_transposed = hidden_states.transpose(-1, -2).contiguous()
-        if hidden_states.shape[-1] > hidden_states.shape[-2]:
-            hidden_states_landscape = hidden_states
-            hidden_states_portrait = hidden_states_transposed
-        else:
-            hidden_states_landscape = hidden_states_transposed
-            hidden_states_portrait = hidden_states
-
-        exported_landscape_1 = torch.export.export(
-            mod=pipeline.transformer,
-            args=call.args,
-            kwargs=call.kwargs | {'hidden_states': hidden_states_landscape},
-            dynamic_shapes=dynamic_shapes,
-        )
-        
-        exported_portrait_2 = torch.export.export(
-            mod=pipeline.transformer_2,
-            args=call.args,
-            kwargs=call.kwargs | {'hidden_states': hidden_states_portrait},
-            dynamic_shapes=dynamic_shapes,
-        )
-
-        compiled_landscape_1 = aoti_compile(exported_landscape_1, INDUCTOR_CONFIGS)
-        compiled_portrait_2 = aoti_compile(exported_portrait_2, INDUCTOR_CONFIGS)
-
-        compiled_landscape_2 = ZeroGPUCompiledModel(compiled_landscape_1.archive_file, compiled_portrait_2.weights)
-        compiled_portrait_1 = ZeroGPUCompiledModel(compiled_portrait_2.archive_file, compiled_landscape_1.weights)
-
-        return (
-            compiled_landscape_1,
-            compiled_landscape_2,
-            compiled_portrait_1,
-            compiled_portrait_2,
-        )
-
+    # --- LoRA 1: lightning (loads into default transformer) ---
+    pipeline.load_lora_weights(
+        "Kijai/WanVideo_comfy",
+        weight_name="Lightx2v/lightx2v_T2V_14B_cfg_step_distill_v2_lora_rank128_bf16.safetensors",
+        adapter_name="lightning",
+    )
+
+    # --- LoRA 2: lightning_2 (loads into transformer_2) ---
+    kwargs_lora = {"load_into_transformer_2": True}
+    pipeline.load_lora_weights(
+        "lightx2v/Wan2.2-Lightning",
+        weight_name="Wan2.2-T2V-A14B-4steps-lora-rank64-Seko-V1/low_noise_model.safetensors",
+        adapter_name="lightning_2",
+        **kwargs_lora,
+    )
+
+    # --- LoRA 3: orbit_shot (the ostris repo you asked for) ---
+    # Load into transformer_2 as well (set load_into_transformer_2 True if this adapter targets transformer_2)
+    pipeline.load_lora_weights(
+        "ostris/wan22_i2v_14b_orbit_shot_lora",
+        weight_name="wan22_14b_i2v_orbit_low_noise.safetensors",
+        adapter_name="orbit_shot",
+        **kwargs_lora,
+    )
+
+    # Register adapters and their relative weights
+    # (adjust adapter_weights to taste; here each is weight 1.0)
+    pipeline.set_adapters(["lightning", "lightning_2", "orbit_shot"], adapter_weights=[1.0, 1.0, 1.0])
+
+    # Fuse each adapter into the correct component with chosen lora_scale:
+    # - lightning -> transformer (boosted by 3x in your original)
+    # - lightning_2 -> transformer_2
+    # - orbit_shot -> transformer_2 (or transformer depending on the LoRA design)
+    pipeline.fuse_lora(adapter_names=["lightning"], lora_scale=3.0, components=["transformer"])
+    pipeline.fuse_lora(adapter_names=["lightning_2"], lora_scale=1.0, components=["transformer_2"])
+    pipeline.fuse_lora(adapter_names=["orbit_shot"], lora_scale=1.0, components=["transformer_2"])
+
+    # After fusing, you can unload LoRA weights to free memory (fused weights remain applied)
+    pipeline.unload_lora_weights()
+
+    # --- then continue with capture_component_call / export / compile logic as you already have ---
+    with capture_component_call(pipeline, 'transformer') as call:
+        pipeline(*args, **kwargs)
+
+    # ... rest of your function unchanged ...
+ 
+   
     quantize_(pipeline.text_encoder, Int8WeightOnlyConfig())
     cl1, cl2, cp1, cp2 = compile_transformer()