Update app_optimized.py

app_optimized.py

@@ -1,3 +1,4 @@
+import os
 import gradio as gr
 import safetensors.torch
 import torchvision.transforms.v2 as transforms
@@ -62,7 +63,7 @@ with torch.no_grad():  # enable image inputs
         new_img_in.proj.bias.copy_(pipe.transformer.x_embedder.proj.bias)
     pipe.transformer.x_embedder = new_img_in
 
-lora_state_dict = safetensors.torch.load_file(lora_path, device="cpu")
+lora_state_dict = safetensors.torch.load_file(lora_path, device="cuda")
 transformer_lora_state_dict = {f'{k.replace("transformer.", "")}': v for k, v in lora_state_dict.items() if k.startswith("transformer.") and "lora" in k}
 pipe.load_lora_into_transformer(transformer_lora_state_dict, transformer=pipe.transformer, adapter_name="i2v", _pipeline=pipe)
 pipe.set_adapters(["i2v"], adapter_weights=[1.0])