Update app.py

app.py

@@ -23,6 +23,7 @@ from transformers import (
     T5Tokenizer
 )
 from accelerate import init_empty_weights
+from accelerate.utils import set_module_tensor_to_device
 from safetensors import safe_open
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -50,8 +51,14 @@ with safe_open(model_file, framework="pt") as f:
         state_dict[key] = f.get_tensor(key)
 
 state_dict = convert_sd3_transformer_checkpoint_to_diffusers(state_dict)
-transformer.load_state_dict(state_dict)
-transformer.to_empty(device=device)
+for key, value in state_dict.items():
+    set_module_tensor_to_device(
+        transformer,
+        key,
+        device,
+        value=value,
+        dtype=torch_dtype
+    )
 
 # Try to keep memory usage down
 del state_dict