update

app.py

@@ -6,8 +6,8 @@ from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor
 from kolors.pipelines import pipeline_stable_diffusion_xl_chatglm_256_ipadapter, pipeline_stable_diffusion_xl_chatglm_256
 from kolors.models.modeling_chatglm import ChatGLMModel
 from kolors.models.tokenization_chatglm import ChatGLMTokenizer
-from kolors.models.unet_2d_condition import UNet2DConditionModel
-from diffusers import AutoencoderKL, EulerDiscreteScheduler
+from kolors.models import unet_2d_condition
+from diffusers import AutoencoderKL, EulerDiscreteScheduler, UNet2DConditionModel
 import gradio as gr
 import numpy as np
 
@@ -17,12 +17,12 @@ 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")
 
-# Load models
 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')
 vae = AutoencoderKL.from_pretrained(f"{ckpt_dir}/vae", revision=None).half().to(device)
 scheduler = EulerDiscreteScheduler.from_pretrained(f"{ckpt_dir}/scheduler")
-unet = UNet2DConditionModel.from_pretrained(f"{ckpt_dir}/unet", revision=None).half().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)
@@ -30,7 +30,7 @@ clip_image_processor = CLIPImageProcessor(size=ip_img_size, crop_size=ip_img_siz
 pipe_t2i = pipeline_stable_diffusion_xl_chatglm_256.StableDiffusionXLPipeline(
     vae=vae,text_encoder=text_encoder, 
     tokenizer=tokenizer, 
-    unet=unet, 
+    unet=unet_t2i, 
     scheduler=scheduler, 
     force_zeros_for_empty_prompt=False
 ).to(device)
@@ -39,7 +39,7 @@ pipe_i2i = pipeline_stable_diffusion_xl_chatglm_256_ipadapter.StableDiffusionXLP
     vae=vae,
     text_encoder=text_encoder,
     tokenizer=tokenizer,
-    unet=unet,
+    unet=unet_i2i,
     scheduler=scheduler,
     image_encoder=image_encoder,
     feature_extractor=clip_image_processor,
@@ -126,7 +126,7 @@ with gr.Blocks(css=css) as demo:
                 )
                 run_button = gr.Button("Run", scale=0)
             with gr.Row():
-                ip_adapter_image = gr.Image(label="IP-Adapter Image", type="pil")
+                ip_adapter_image = gr.Image(label="IP-Adapter Image (optional)", type="pil")
             with gr.Accordion("Advanced Settings", open=False):
                 negative_prompt = gr.Text(
                     label="Negative prompt",