Update app.py

app.py

@@ -52,20 +52,18 @@ def face_bbox_to_square(bbox):
     return [l0, t0, r0, b0]
 
 
-ckpt_dir = f'weights/Kolors'
-text_encoder = ChatGLMModel.from_pretrained(
-    f'{ckpt_dir}/text_encoder').to(dtype=torch.bfloat16)
-tokenizer = ChatGLMTokenizer.from_pretrained(f'{ckpt_dir}/text_encoder')
-vae = AutoencoderKL.from_pretrained(f"{ckpt_dir}/vae", revision=None).to(dtype=torch.bfloat16)
-scheduler = EulerDiscreteScheduler.from_pretrained(f"{ckpt_dir}/scheduler")
-unet = UNet2DConditionModel.from_pretrained(f"{ckpt_dir}/unet", revision=None).to(dtype=torch.bfloat16)
-
-control_path = f'weights/Kolors-Controlnet-Pose-Tryon'
+text_encoder = ChatGLMModel.from_pretrained("Kwai-Kolors/Kolors",subfolder="text_encoder").to(dtype=torch.bfloat16)
+tokenizer = ChatGLMTokenizer.from_pretrained("Kwai-Kolors/Kolors",subfolder="text_encoder")
+vae = AutoencoderKL.from_pretrained("Kwai-Kolors/Kolors",subfolder="vae",, revision=None).to(dtype=torch.bfloat16)
+scheduler = EulerDiscreteScheduler.from_pretrained("Kwai-Kolors/Kolors",subfolder="scheduler")
+unet = UNet2DConditionModel.from_pretrained("Kwai-Kolors/Kolors",subfolder="unet", revision=None).to(dtype=torch.bfloat16)
+
+control_path = "haowu11/Kolors-Controlnet-Pose-Tryon"
 controlnet = ControlNetModel.from_pretrained( control_path , revision=None).to(dtype=torch.bfloat16)
 
 face_info_generator = FaceInfoGenerator(root_dir = "./")
 
-clip_image_encoder = CLIPVisionModelWithProjection.from_pretrained(f'weights/Kolors-IP-Adapter-FaceID-Plus/clip-vit-large-patch14-336', ignore_mismatched_sizes=True)
+clip_image_encoder = CLIPVisionModelWithProjection.from_pretrained("Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus",subfolder="clip-vit-large-patch14-336", ignore_mismatched_sizes=True)
 clip_image_encoder.to('cuda')
 clip_image_processor = CLIPImageProcessor(size = 336, crop_size = 336)
 
@@ -85,7 +83,7 @@ pipe = StableDiffusionXLControlNetImg2ImgPipeline(
 if hasattr(pipe.unet, 'encoder_hid_proj'):
     pipe.unet.text_encoder_hid_proj = pipe.unet.encoder_hid_proj
 ip_scale = 0.5
-pipe.load_ip_adapter_faceid_plus(f'weights/Kolors-IP-Adapter-FaceID-Plus/ipa-faceid-plus.bin', device = 'cuda')
+pipe.load_ip_adapter_faceid_plus(f'ipa-faceid-plus.bin', device = 'cuda')
 pipe.set_face_fidelity_scale(ip_scale)
 pipe = pipe.to("cuda")
 pipe.enable_model_cpu_offload()
@@ -205,5 +203,5 @@ with block:
         run_button = gr.Button(value="Run")
     ips = [face_img,pose_img, garm_img, prompt,negative_prompt, n_samples, n_steps, seed]
     run_button.click(fn=infer, inputs=ips, outputs=[result_gallery])
-
-block.launch(server_name='0.0.0.0', server_port=7865)
+if __name__ == "__main__":
+    block.launch(server_name='0.0.0.0')