update

app.py

@@ -17,24 +17,24 @@ ckpt_IPA_dir = '/home/lixiang46/Kolors/weights/Kolors-IP-Adapter-Plus'
 # ckpt_dir = snapshot_download(repo_id="Kwai-Kolors/Kolors")
 # ckpt_IPA_dir = snapshot_download(repo_id="Kwai-Kolors/Kolors-IP-Adapter-Plus")
 
-text_encoder = ChatGLMModel.from_pretrained(f'{ckpt_dir}/text_encoder', torch_dtype=torch.float16).half()
-tokenizer = ChatGLMTokenizer.from_pretrained(f'{ckpt_dir}/text_encoder')
-vae = AutoencoderKL.from_pretrained(f"{ckpt_dir}/vae", revision=None).half()
-scheduler = EulerDiscreteScheduler.from_pretrained(f"{ckpt_dir}/scheduler")
-# unet_t2i = UNet2DConditionModel.from_pretrained(f"{ckpt_dir}/unet", revision=None).half()
-unet_i2i = unet_2d_condition.UNet2DConditionModel.from_pretrained(f"{ckpt_dir}/unet", revision=None).half()
+text_encoder = ChatGLMModel.from_pretrained(f'{ckpt_dir}/text_encoder', torch_dtype=torch.float16).half().to(device)
+tokenizer = ChatGLMTokenizer.from_pretrained(f'{ckpt_dir}/text_encoder').to(device)
+vae = AutoencoderKL.from_pretrained(f"{ckpt_dir}/vae", revision=None).half().to(device)
+scheduler = EulerDiscreteScheduler.from_pretrained(f"{ckpt_dir}/scheduler").to(device)
+unet_t2i = UNet2DConditionModel.from_pretrained(f"{ckpt_dir}/unet", revision=None).half().to(device)
+unet_i2i = unet_2d_condition.UNet2DConditionModel.from_pretrained(f"{ckpt_dir}/unet", revision=None).half().to(device)
 image_encoder = CLIPVisionModelWithProjection.from_pretrained(f'{ckpt_IPA_dir}/image_encoder',ignore_mismatched_sizes=True).to(dtype=torch.float16, device=device)
 ip_img_size = 336
 clip_image_processor = CLIPImageProcessor(size=ip_img_size, crop_size=ip_img_size)
 
-# pipe_t2i = pipeline_stable_diffusion_xl_chatglm_256.StableDiffusionXLPipeline(
-#     vae=vae,
-#     text_encoder=text_encoder, 
-#     tokenizer=tokenizer, 
-#     unet=unet_t2i, 
-#     scheduler=scheduler, 
-#     force_zeros_for_empty_prompt=False
-# ).to(device)
+pipe_t2i = pipeline_stable_diffusion_xl_chatglm_256.StableDiffusionXLPipeline(
+    vae=vae,
+    text_encoder=text_encoder, 
+    tokenizer=tokenizer, 
+    unet=unet_t2i, 
+    scheduler=scheduler, 
+    force_zeros_for_empty_prompt=False
+).to(device)
 
 pipe_i2i = pipeline_stable_diffusion_xl_chatglm_256_ipadapter.StableDiffusionXLPipeline(
     vae=vae,
@@ -61,7 +61,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
     generator = torch.Generator().manual_seed(seed)
 
     if ip_adapter_image is None:
-        image = pipe_i2i(
+        image = pipe_t2i(
             prompt = prompt, 
             negative_prompt = negative_prompt,
             guidance_scale = guidance_scale,