22

app.py

@@ -37,11 +37,8 @@ transformer = ZImageControlTransformer2DModel.from_pretrained(
 
 if TRANSFORMER_LOCAL is not None:
     print(f"From checkpoint: {TRANSFORMER_LOCAL}")
-    if TRANSFORMER_LOCAL.endswith("safetensors"):
-        from safetensors.torch import load_file, safe_open
-        state_dict = load_file(TRANSFORMER_LOCAL)
-    else:
-        state_dict = torch.load(TRANSFORMER_LOCAL, map_location="cpu")
+    from safetensors.torch import load_file, safe_open
+    state_dict = load_file(TRANSFORMER_LOCAL)
     state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
 
     m, u = transformer.load_state_dict(state_dict, strict=False)