clean up

gradio_ui.py

@@ -1,88 +1,107 @@
-import spaces
-import PIL
+import os
 import torch
-import subprocess
+import PIL
 import gradio as gr
-import os
 
 from typing import Optional
 from accelerate import Accelerator
 from diffusers import (
     AutoencoderKL,
-    StableDiffusionXLControlNetPipeline, 
+    StableDiffusionXLControlNetPipeline,
     ControlNetModel,
     UNet2DConditionModel,
 )
 from transformers import (
-    BlipProcessor, BlipForConditionalGeneration, 
-    VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
+    BlipProcessor, BlipForConditionalGeneration,
 )
-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 hf_hub_download, snapshot_download
 
-from huggingface_hub import snapshot_download
+import spaces
 
-# Download colorization models
+
+# ========== Initialization ==========
+
+# Ensure required directories exist
 os.makedirs("sdxl_light_caption_output", exist_ok=True)
-os.makedirs("sdxl_light_custom_caption_output", exist_ok=True)
 
+# Download controlnet model snapshot
 snapshot_download(
-    repo_id = 'nickpai/sdxl_light_caption_output',
-    local_dir = 'sdxl_light_caption_output'
+    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'
+# Device and precision setup
+accelerator = Accelerator(mixed_precision="fp16")
+weight_dtype = torch.float16 if accelerator.mixed_precision == "fp16" else torch.float32
+device = accelerator.device
+
+print(f"[INFO] Accelerator device: {device}")
+
+# ========== Models ==========
+
+# Pretrained paths
+base_model_path = "stabilityai/stable-diffusion-xl-base-1.0"
+safetensors_ckpt = "sdxl_lightning_8step_unet.safetensors"
+controlnet_path = "sdxl_light_caption_output/checkpoint-30000/controlnet"
+
+# Load diffusion components
+vae = AutoencoderKL.from_pretrained(base_model_path, subfolder="vae")
+unet = UNet2DConditionModel.from_config(base_model_path, subfolder="unet")
+unet.load_state_dict(load_file(hf_hub_download("ByteDance/SDXL-Lightning", safetensors_ckpt)))
+
+controlnet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=weight_dtype)
+
+pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+    base_model_path, vae=vae, unet=unet, controlnet=controlnet
 )
+pipe.to(device, dtype=weight_dtype)
+pipe.safety_checker = None
+
+# Load BLIP captioning model
+caption_model_name = "blip-image-captioning-large"
+processor = BlipProcessor.from_pretrained(f"Salesforce/{caption_model_name}")
+caption_model = BlipForConditionalGeneration.from_pretrained(
+    f"Salesforce/{caption_model_name}", torch_dtype=weight_dtype
+).to(device)
 
+# ========== Utility Functions ==========
 
 def apply_color(image: PIL.Image.Image, color_map: PIL.Image.Image) -> PIL.Image.Image:
-    # Convert input images to LAB color space
+    # Convert to LAB color space
     image_lab = image.convert('LAB')
     color_map_lab = color_map.convert('LAB')
 
-    # Split LAB channels
-    l, a , b = image_lab.split()
+    # Extract and merge LAB channels
+    l, _, _ = image_lab.split()
     _, a_map, b_map = color_map_lab.split()
-
-    # Merge LAB channels with color map
     merged_lab = PIL.Image.merge('LAB', (l, a_map, b_map))
 
-    # Convert merged LAB image back to RGB color space
-    result_rgb = merged_lab.convert('RGB')
-    return result_rgb
-
-def remove_unlikely_words(prompt: str) -> str:
-    """
-    Removes unlikely words from a prompt.
+    return merged_lab.convert('RGB')
 
-    Args:
-        prompt: The text prompt to be cleaned.
 
-    Returns:
-        The cleaned prompt with unlikely words removed.
-    """
+def remove_unlikely_words(prompt: str) -> str:
+    """Removes predefined unlikely phrases from prompt text."""
     unlikely_words = []
 
-    a1_list = [f'{i}s' for i in range(1900, 2000)]
-    a2_list = [f'{i}' for i in range(1900, 2000)]
-    a3_list = [f'year {i}' for i in range(1900, 2000)]
-    a4_list = [f'circa {i}' for i in range(1900, 2000)]
-    b1_list = [f"{year[0]} {year[1]} {year[2]} {year[3]} s" for year in a1_list]
-    b2_list = [f"{year[0]} {year[1]} {year[2]} {year[3]}" for year in a1_list]
-    b3_list = [f"year {year[0]} {year[1]} {year[2]} {year[3]}" for year in a1_list]
-    b4_list = [f"circa {year[0]} {year[1]} {year[2]} {year[3]}" for year in a1_list]
+    a1 = [f'{i}s' for i in range(1900, 2000)]
+    a2 = [f'{i}' for i in range(1900, 2000)]
+    a3 = [f'year {i}' for i in range(1900, 2000)]
+    a4 = [f'circa {i}' for i in range(1900, 2000)]
+
+    b1 = [f"{y[0]} {y[1]} {y[2]} {y[3]} s" for y in a1]
+    b2 = [f"{y[0]} {y[1]} {y[2]} {y[3]}" for y in a1]
+    b3 = [f"year {y[0]} {y[1]} {y[2]} {y[3]}" for y in a1]
+    b4 = [f"circa {y[0]} {y[1]} {y[2]} {y[3]}" for y in a1]
 
-    words_list = [
+    manual = [  # same list as your original words_list
         "black and white,", "black and white", "black & white,", "black & white", "circa", 
         "balck and white,", "monochrome,", "black-and-white,", "black-and-white photography,", 
         "black - and - white photography,", "monochrome bw,", "black white,", "black an white,",
         "grainy footage,", "grainy footage", "grainy photo,", "grainy photo", "b&w photo",
         "back and white", "back and white,", "monochrome contrast", "monochrome", "grainy",
         "grainy photograph,", "grainy photograph", "low contrast,", "low contrast", "b & w",
-        "grainy black-and-white photo,", "bw", "bw,",  "grainy black-and-white photo",
+        "grainy black-and-white photo,", "bw", "bw,", "grainy black-and-white photo",
         "b & w,", "b&w,", "b&w!,", "b&w", "black - and - white,", "bw photo,", "grainy  photo,",
         "black-and-white photo,", "black-and-white photo", "black - and - white photography",
         "b&w photo,", "monochromatic photo,", "grainy monochrome photo,", "monochromatic",
@@ -94,287 +113,85 @@ def remove_unlikely_words(prompt: str) -> str:
         "historical photo", "historical setting,",
         "historic photo,", "historic", "desaturated!!,", "desaturated!,", "desaturated,", "desaturated", 
         "taken in", "shot on leica", "shot on leica sl2", "sl2",
-        "taken with a leica camera", "taken with a leica camera", "leica sl2", "leica", "setting", 
+        "taken with a leica camera", "leica sl2", "leica", "setting", 
         "overcast day", "overcast weather", "slight overcast", "overcast", 
         "picture taken in", "photo taken in", 
         ", photo", ",  photo", ",   photo", ",    photo", ", photograph",
         ",,", ",,,", ",,,,", " ,", "  ,", "   ,", "    ,", 
     ]
 
-    unlikely_words.extend(a1_list)
-    unlikely_words.extend(a2_list)
-    unlikely_words.extend(a3_list)
-    unlikely_words.extend(a4_list)
-    unlikely_words.extend(b1_list)
-    unlikely_words.extend(b2_list)
-    unlikely_words.extend(b3_list)
-    unlikely_words.extend(b4_list)
-    unlikely_words.extend(words_list)
-    
+    unlikely_words.extend(a1 + a2 + a3 + a4 + b1 + b2 + b3 + b4 + manual)
+
     for word in unlikely_words:
         prompt = prompt.replace(word, "")
     return prompt
 
-def blip_image_captioning(image: PIL.Image.Image,
-                          model_backbone: str,
-                          weight_dtype: type,
-                          device: str,
-                          conditional: bool) -> str:
-    # https://huggingface.co/Salesforce/blip-image-captioning-large
-    # https://huggingface.co/Salesforce/blip-image-captioning-base
-    if weight_dtype == torch.bfloat16: # in case model might not accept bfloat16 data type
-        weight_dtype = torch.float16
-
-    processor = BlipProcessor.from_pretrained(f"Salesforce/{model_backbone}")
-    model = BlipForConditionalGeneration.from_pretrained(
-         f"Salesforce/{model_backbone}", torch_dtype=weight_dtype).to(device)
-    
-    valid_backbones = ["blip-image-captioning-large", "blip-image-captioning-base"]
-    if model_backbone not in valid_backbones:
-         raise ValueError(f"Invalid model backbone '{model_backbone}'. \
-                          Valid options are: {', '.join(valid_backbones)}")
-
-    if conditional:
-        text = "a photography of"
-        inputs = processor(image, text, return_tensors="pt").to(device, weight_dtype)
-    else:
-        inputs = processor(image, return_tensors="pt").to(device)
-    out = model.generate(**inputs)
-    caption = processor.decode(out[0], skip_special_tokens=True)
-    return caption
-
-# def vit_gpt2_image_captioning(image: PIL.Image.Image, device: str) -> str:
-#     # https://huggingface.co/nlpconnect/vit-gpt2-image-captioning
-#     model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning").to(device)
-#     feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
-#     tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
-
-#     max_length = 16
-#     num_beams = 4
-#     gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
-
-#     pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
-#     pixel_values = pixel_values.to(device)
-
-#     output_ids = model.generate(pixel_values, **gen_kwargs)
-
-#     preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
-#     caption = [pred.strip() for pred in preds]
-
-#     return caption[0]
-
-# def clip_image_captioning(image: PIL.Image.Image,
-#                           clip_model_name: str,
-#                           device: str) -> str:
-#     # validate clip model name
-#     models = list_clip_models()
-#     if clip_model_name not in models:
-#         raise ValueError(f"Could not find CLIP model {clip_model_name}! \
-#                          Available models: {models}")
-#     config = Config(device=device, clip_model_name=clip_model_name)
-#     config.apply_low_vram_defaults()
-#     ci = Interrogator(config)
-#     caption = ci.interrogate(image)
-#     return caption
-
-# Define a function to process the image with the loaded model
+
+def get_image_paths(folder_path: str) -> list:
+    return [[os.path.join(folder_path, f)] for f in os.listdir(folder_path)
+            if f.lower().endswith((".jpg", ".png"))]
+
+
 @spaces.GPU
-def process_image(image_path: str,
-                  controlnet_model_name_or_path: str,
-                  caption_model_name: str,
+def process_image(image_path: str,                 
                   positive_prompt: Optional[str],
                   negative_prompt: Optional[str],
-                  seed: int,
-                  num_inference_steps: int,
-                  mixed_precision: str,
-                  pretrained_model_name_or_path: str,
-                  pretrained_vae_model_name_or_path: Optional[str],
-                  revision: Optional[str],
-                  variant: Optional[str],
-                  repo: str,
-                  ckpt: str,) -> PIL.Image.Image:
-    # Seed
-    generator = torch.manual_seed(seed)
-
-    # Accelerator Setting
-    accelerator = Accelerator(
-        mixed_precision=mixed_precision,
-        cpu=False
-    )
-
-    print(f"Accelerator device: {accelerator.device}")
-    
-    weight_dtype = torch.float32
-    if accelerator.mixed_precision == "fp16":
-        weight_dtype = torch.float16
-    elif accelerator.mixed_precision == "bf16":
-        weight_dtype = torch.bfloat16
-
-    vae_path = (
-        pretrained_model_name_or_path
-        if pretrained_vae_model_name_or_path is None
-        else pretrained_vae_model_name_or_path
-    )
-    vae = AutoencoderKL.from_pretrained(
-        vae_path,
-        subfolder="vae" if pretrained_vae_model_name_or_path is None else None,
-        revision=revision,
-        variant=variant,
-    )
-    unet = UNet2DConditionModel.from_config(
-        pretrained_model_name_or_path, 
-        subfolder="unet", 
-        revision=revision, 
-        variant=variant,
-    )
-    unet.load_state_dict(load_file(hf_hub_download(repo, ckpt)))
-
-    # Move vae, unet and text_encoder to device and cast to weight_dtype
-    # The VAE is in float32 to avoid NaN losses.
-    if pretrained_vae_model_name_or_path is not None:
-        vae.to(accelerator.device, dtype=weight_dtype)
-    else:
-        vae.to(accelerator.device, dtype=torch.float32)
-    unet.to(accelerator.device, dtype=weight_dtype)
-
-    controlnet = ControlNetModel.from_pretrained(controlnet_model_name_or_path, torch_dtype=weight_dtype)
-    pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
-        pretrained_model_name_or_path,
-        vae=vae,
-        unet=unet,
-        controlnet=controlnet, 
-    )
-    pipe.to(accelerator.device, dtype=weight_dtype)
+                  seed: int) -> tuple[PIL.Image.Image, str]:
 
+    torch.manual_seed(seed)
     image = PIL.Image.open(image_path)
-
-    # Prepare everything with our `accelerator`.
-    pipe, image = accelerator.prepare(pipe, image)
-    pipe.safety_checker = None
-
-    # Convert image into grayscale
     original_size = image.size
     control_image = image.convert("L").convert("RGB").resize((512, 512))
-    
+
     # Image captioning
-    if caption_model_name == "blip-image-captioning-large" or "blip-image-captioning-base":
-        caption = blip_image_captioning(control_image, caption_model_name, 
-                                        weight_dtype, accelerator.device, conditional=True)
-    # elif caption_model_name == "ViT-L-14/openai" or "ViT-H-14/laion2b_s32b_b79k":
-    #     caption = clip_image_captioning(control_image, caption_model_name, accelerator.device)
-    # elif caption_model_name == "vit-gpt2-image-captioning":
-    #     caption = vit_gpt2_image_captioning(control_image, accelerator.device)
+    input_text = "a photography of"
+    inputs = processor(image, input_text, return_tensors="pt").to(device, dtype=weight_dtype)
+    caption_ids = caption_model.generate(**inputs)
+    caption = processor.decode(caption_ids[0], skip_special_tokens=True)
     caption = remove_unlikely_words(caption)
 
-    # Combine positive prompt and captioning result
-    prompt = [positive_prompt + ", " + caption]
-
-    # Image colorization
-    image = pipe(prompt=prompt, 
-                 negative_prompt=negative_prompt, 
-                 num_inference_steps=num_inference_steps, 
-                 generator=generator, 
-                 image=control_image).images[0]
-    
-    # Apply color mapping
-    result_image = apply_color(control_image, image)
-    result_image = result_image.resize(original_size)
-    return result_image, caption
-
-# Define the image gallery based on folder path
-def get_image_paths(folder_path):
-  import os
-  image_paths = []
-  for filename in os.listdir(folder_path):
-    if filename.endswith(".jpg") or filename.endswith(".png"):
-      image_paths.append([os.path.join(folder_path, filename)])
-  return image_paths
-
-# Create the Gradio interface
+    # Inference
+    final_prompt = [f"{positive_prompt}, {caption}"]
+    result = pipe(prompt=final_prompt,
+                  negative_prompt=negative_prompt,
+                  num_inference_steps=8,
+                  generator=torch.manual_seed(seed),
+                  image=control_image)
+
+    colorized = apply_color(control_image, result.images[0]).resize(original_size)
+    return colorized, caption
+
+
+# ========== Gradio UI ==========
+
 def create_interface():
-    controlnet_model_dict = {
-       "sdxl-light-caption-30000": "sdxl_light_caption_output/checkpoint-30000/controlnet",
-       "sdxl-light-custom-caption-30000": "sdxl_light_custom_caption_output/checkpoint-30000/controlnet",
-    }
-    images = get_image_paths("example/legacy_images")  # Replace with your folder path
+    examples = get_image_paths("example/legacy_images")
 
-    interface = gr.Interface(
+    return gr.Interface(
         fn=process_image,
         inputs=[
-            gr.Image(label="Upload image", 
-                     value="example/legacy_images/Hollywood-Sign.jpg",
-                     type='filepath'),
-            gr.Dropdown(choices=[controlnet_model_dict[key] for key in controlnet_model_dict], 
-                        value=controlnet_model_dict["sdxl-light-caption-30000"],
-                        label="Select ControlNet Model"),
-            gr.Dropdown(choices=["blip-image-captioning-large",
-                                 "blip-image-captioning-base",], 
-                        value="blip-image-captioning-large",
-                        label="Select Image Captioning Model"),
-            gr.Textbox(label="Positive Prompt", placeholder="Text for positive prompt"),
-            gr.Textbox(value="low quality, bad quality, low contrast, black and white, bw, monochrome, grainy, blurry, historical, restored, desaturate",
-                       label="Negative Prompt", placeholder="Text for negative prompt"),
+            gr.Image(label="Upload Image", type='filepath',
+                     value="example/legacy_images/Hollywood-Sign.jpg"),
+            gr.Textbox(label="Positive Prompt", placeholder="Enter details to enhance the caption"),
+            gr.Textbox(label="Negative Prompt", value="low quality, bad quality, low contrast, black and white, bw, monochrome, grainy, blurry, historical, restored, desaturate"),
         ],
         outputs=[
-            gr.Image(label="Colorized image", 
-                     value="example/UUColor_results/Hollywood-Sign.jpeg",
-                     format="jpeg"),
-            gr.Textbox(label="Captioning Result", show_copy_button=True)
-        ],
-        examples=images,
-        additional_inputs=[
-            # gr.Radio(choices=["Original", "Square"], value="Original", 
-            #          label="Output resolution"),
-            # gr.Slider(minimum=128, maximum=512, value=256, step=128,
-            #           label="Height & Width", 
-            #           info='Only effect if select "Square" output resolution'),
-            gr.Slider(0, 1000, 123, label="Seed"),
-            gr.Radio(choices=[1, 2, 4, 8],
-                     value=8,
-                     label="Inference Steps",
-                     info="1-step, 2-step, 4-step, or 8-step distilled models"),
-            gr.Radio(choices=["no", "fp16", "bf16"],
-                     value="fp16",
-                     label="Mixed Precision",
-                     info="Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16)."),
-            gr.Dropdown(choices=["stabilityai/stable-diffusion-xl-base-1.0"],
-                        value="stabilityai/stable-diffusion-xl-base-1.0",
-                        label="Base Model",
-                        info="Path to pretrained model or model identifier from huggingface.co/models."),
-            gr.Dropdown(choices=["None"],
-                        value=None,
-                        label="VAE Model",
-                        info="Path to an improved VAE to stabilize training. For more details check out: https://github.com/huggingface/diffusers/pull/4038."),
-            gr.Dropdown(choices=["None"],
-                        value=None,
-                        label="Varient",
-                        info="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16"),
-            gr.Dropdown(choices=["None"],
-                        value=None,
-                        label="Revision",
-                        info="Revision of pretrained model identifier from huggingface.co/models."),
-            gr.Dropdown(choices=["ByteDance/SDXL-Lightning"],
-                        value="ByteDance/SDXL-Lightning",
-                        label="Repository",
-                        info="Repository from huggingface.co"),
-            gr.Dropdown(choices=["sdxl_lightning_1step_unet.safetensors",
-                                 "sdxl_lightning_2step_unet.safetensors",
-                                 "sdxl_lightning_4step_unet.safetensors",
-                                 "sdxl_lightning_8step_unet.safetensors"],
-                        value="sdxl_lightning_8step_unet.safetensors",
-                        label="Checkpoint",
-                        info="Available checkpoints from the repository. Caution! Checkpoint's 'N'step must match with inference steps"),
+            gr.Image(label="Colorized Image", format="jpeg",
+                     value="example/UUColor_results/Hollywood-Sign.jpeg"),
+            gr.Textbox(label="Caption", show_copy_button=True)
         ],
+        examples=examples,
+        additional_inputs=[gr.Slider(0, 1000, 123, label="Seed")],
         title="Text-Guided Image Colorization",
-        description="Upload an image and select a model to colorize it.",
+        description="Upload a grayscale image and generate a color version guided by automatic captioning.",
         cache_examples=False
     )
-    return interface
+
 
 def main():
-    # Launch the Gradio interface
     interface = create_interface()
     interface.launch(ssr_mode=False)
 
+
 if __name__ == "__main__":
-   main()
+    main()