Update gradio_ui.py

gradio_ui.py

@@ -2,6 +2,7 @@ import PIL
 import torch
 import subprocess
 import gradio as gr
+import os
 
 from typing import Optional
 from accelerate import Accelerator
@@ -19,6 +20,16 @@ from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 from clip_interrogator import Interrogator, Config, list_clip_models
 
+from huggingface_hub import snapshot_download
+
+# Ensure 'checkpoint' directory exists
+os.makedirs("sdxl_light_caption_output", exist_ok=True)
+
+snapshot_download(
+    repo_id = 'nickpai/sdxl_light_caption_output',
+    local_dir = 'sdxl_light_caption_output'
+)
+
 def apply_color(image: PIL.Image.Image, color_map: PIL.Image.Image) -> PIL.Image.Image:
     # Convert input images to LAB color space
     image_lab = image.convert('LAB')