Update app.py

app.py

@@ -1,4 +1,5 @@
 import os, torch, random
+from huggingface_hub import snapshot_download
 from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256 import StableDiffusionXLPipeline
 from kolors.models.modeling_chatglm import ChatGLMModel
 from kolors.models.tokenization_chatglm import ChatGLMTokenizer
@@ -6,14 +7,18 @@ from diffusers import UNet2DConditionModel, AutoencoderKL
 from diffusers import EulerDiscreteScheduler
 import gradio as gr
 
-ckpt_dir = f"Kwai-Kolors/Kolors"
+# Download the model files
+ckpt_dir = snapshot_download(repo_id="Kwai-Kolors/Kolors")
+
+# Load the models
 text_encoder = ChatGLMModel.from_pretrained(
-    f'{ckpt_dir}/text_encoder',
+    os.path.join(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 = UNet2DConditionModel.from_pretrained(f"{ckpt_dir}/unet", revision=None).half()
+tokenizer = ChatGLMTokenizer.from_pretrained(os.path.join(ckpt_dir, 'text_encoder'))
+vae = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir, "vae"), revision=None).half()
+scheduler = EulerDiscreteScheduler.from_pretrained(os.path.join(ckpt_dir, "scheduler"))
+unet = UNet2DConditionModel.from_pretrained(os.path.join(ckpt_dir, "unet"), revision=None).half()
+
 pipe = StableDiffusionXLPipeline(
         vae=vae,
         text_encoder=text_encoder,
@@ -55,4 +60,4 @@ iface = gr.Interface(
     description="Generate images using the Kolors Stable Diffusion XL model."
 )
 
-iface.launch()
+iface.launch(debug=True)