Upload app.py

app.py

@@ -41,50 +41,47 @@ from insightface.utils import face_align
 
 # device = 'cuda:2' if torch.cuda.is_available() else 'cpu'
 
-parser = argparse.ArgumentParser(description='ReferenceAdapter diffusion')
-parser.add_argument('--if_resampler', type=bool, default=True)
+parser = argparse.ArgumentParser(description='IMAGDressing-v1')
+# parser.add_argument('--if_resampler', type=bool, default=True)
 parser.add_argument('--if_ipa', type=bool, default=True)
 parser.add_argument('--if_control', type=bool, default=True)
 
 # parser.add_argument('--pretrained_model_name_or_path',
 #                     default="./ckpt/Realistic_Vision_V4.0_noVAE",
 #                     type=str)
-parser.add_argument('--ip_ckpt',
-                    default="./ckpt/ip-adapter-faceid-plus_sd15.bin",
-                    type=str)
-parser.add_argument('--pretrained_image_encoder_path',
-                    default="./ckpt/image_encoder/",
-                    type=str)
-parser.add_argument('--pretrained_vae_model_path',
-                    default="./ckpt/sd-vae-ft-mse/",
-                    type=str)
-parser.add_argument('--model_ckpt',
-                    default="./ckpt/IMAGDressing-v1_512.pt",
-                    type=str)
-parser.add_argument('--output_path', type=str, default="./output_ipa_control_resampler")
-# parser.add_argument('--device', type=str, default="cuda:0")
+# parser.add_argument('--ip_ckpt',
+#                     default="./ckpt/ip-adapter-faceid-plus_sd15.bin",
+#                     type=str)
+# parser.add_argument('--pretrained_image_encoder_path',
+#                     default="./ckpt/image_encoder/",
+#                     type=str)
+# parser.add_argument('--pretrained_vae_model_path',
+#                     default="./ckpt/sd-vae-ft-mse/",
+#                     type=str)
+# parser.add_argument('--model_ckpt',
+#                     default="./ckpt/IMAGDressing-v1_512.pt",
+#                     type=str)
+# parser.add_argument('--output_path', type=str, default="./output_ipa_control_resampler")
+# # parser.add_argument('--device', type=str, default="cuda:0")
 args = parser.parse_args()
 
 # svae path
-output_path = args.output_path
+# output_path = args.output_path
+#
+# if not os.path.exists(output_path):
+#     os.makedirs(output_path)
 
-if not os.path.exists(output_path):
-    os.makedirs(output_path)
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-args.device = device
+args.device = "cuda"
 
 base_path = 'feishen29/IMAGDressing-v1'
 
 generator = torch.Generator(device=args.device).manual_seed(42)
-vae = AutoencoderKL.from_pretrained(args.pretrained_vae_model_path).to(dtype=torch.float16, device=args.device)
+vae = AutoencoderKL.from_pretrained('./ckpt/sd-vae-ft-mse/').to(dtype=torch.float16, device=args.device)
 tokenizer = CLIPTokenizer.from_pretrained("./ckpt/tokenizer")
-text_encoder = CLIPTextModel.from_pretrained("./ckpt/text_encoder").to(
-    dtype=torch.float16, device=args.device)
-image_encoder = CLIPVisionModelWithProjection.from_pretrained(args.pretrained_image_encoder_path).to(
-    dtype=torch.float16, device=args.device)
-unet = UNet2DConditionModel.from_pretrained("./ckpt/unet").to(
-    dtype=torch.float16,device=args.device)
+text_encoder = CLIPTextModel.from_pretrained("./ckpt/text_encoder").to(dtype=torch.float16, device=args.device)
+image_encoder = CLIPVisionModelWithProjection.from_pretrained('./ckpt/image_encoder/').to(dtype=torch.float16, device=args.device)
+unet = UNet2DConditionModel.from_pretrained("./ckpt/unet").to(dtype=torch.float16,device=args.device)
 
 # image_face_fusion = pipeline('face_fusion_torch', model='damo/cv_unet_face_fusion_torch', model_revision='v1.0.3')
 
@@ -136,7 +133,7 @@ ref_unet.set_attn_processor(
     {name: CacheAttnProcessor2_0() for name in ref_unet.attn_processors.keys()})  # set cache
 
 # weights load
-model_sd = torch.load(args.model_ckpt, map_location="cpu")["module"]
+model_sd = torch.load('./ckpt/IMAGDressing-v1_512.pt', map_location="cpu")["module"]
 
 ref_unet_dict = {}
 unet_dict = {}
@@ -257,7 +254,7 @@ def dress_process(garm_img, face_img, pose_img, prompt, cloth_guidance_scale, ca
     # noise_scheduler = UniPCMultistepScheduler.from_config(args.pretrained_model_name_or_path, subfolder="scheduler")
     pipe = PipIpaControlNet(unet=unet, reference_unet=ref_unet, vae=vae, tokenizer=tokenizer,
                             text_encoder=text_encoder, image_encoder=image_encoder,
-                            ip_ckpt=args.ip_ckpt,
+                            ip_ckpt='./ckpt/ip-adapter-faceid-plus_sd15.bin',
                             ImgProj=image_proj, controlnet=control_net_openpose,
                             scheduler=noise_scheduler,
                             safety_checker=StableDiffusionSafetyChecker,