Update gradio_ui.py

gradio_ui.py

@@ -22,14 +22,21 @@ from clip_interrogator import Interrogator, Config, list_clip_models
 
 from huggingface_hub import snapshot_download
 
-# Ensure 'checkpoint' directory exists
+# Download colorization models
 os.makedirs("sdxl_light_caption_output", exist_ok=True)
+os.makedirs("sdxl_light_custom_caption_output", exist_ok=True)
 
 snapshot_download(
     repo_id = 'nickpai/sdxl_light_caption_output',
     local_dir = 'sdxl_light_caption_output'
 )
 
+snapshot_download(
+    repo_id = 'nickpai/sdxl_light_custom_caption_output',
+    local_dir = 'sdxl_light_custom_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')