22

app.py

@@ -35,25 +35,25 @@ has_merged = repo_utils.check_dir_exist(TRANSFORMER_MERGED)
 # load transformer
 config = OmegaConf.load(TRANSFORMER_CONFIG)
 
-if not has_merged:
-    print('load transformer from base')
-    transformer = ZImageControlTransformer2DModel.from_pretrained(
-        MODEL_LOCAL,
-        subfolder="transformer",
-        transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
-    ).to("cuda", torch.bfloat16)
-    print('load state_dict')
-    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)
-    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
-    transformer.save_pretrained(TRANSFORMER_MERGED)
-else:
-    print('load transformer from merged to bypass calculation')
-    transformer = ZImageControlTransformer2DModel.from_pretrained(
-        TRANSFORMER_MERGED,
-        transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
-    ).to("cuda", torch.bfloat16)
+# if not has_merged:
+print('load transformer from base')
+transformer = ZImageControlTransformer2DModel.from_pretrained(
+    MODEL_LOCAL,
+    subfolder="transformer",
+    transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
+).to("cuda", torch.bfloat16)
+print('load state_dict')
+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)
+print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
+transformer.save_pretrained(TRANSFORMER_MERGED)
+# else:
+#     print('load transformer from merged to bypass calculation')
+#     transformer = ZImageControlTransformer2DModel.from_pretrained(
+#         TRANSFORMER_MERGED,
+#         transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
+#     ).to("cuda", torch.bfloat16)
 
 print("transformer ready.")