add new model

app.py

@@ -5,13 +5,73 @@ from torchvision import transforms
 from matplotlib import pyplot as plt
 import gradio as gr
 
+import transformers
+transformers.utils.move_cache()
+
+from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
+from transformers import BlipProcessor, BlipForConditionalGeneration
+from accelerate import Accelerator
+
+import warnings
+warnings.filterwarnings("ignore") 
+
+
 from models import MainModel, UNetAuto, Autoencoder  
 from utils import lab_to_rgb, build_res_unet, build_mobilenet_unet  # Utility to convert LAB to RGB
+from stable import blip_image_captioning, apply_color
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Hàm load models
-def load_unet_model(auto_model_path):
+
+# Stable diffusion
+
+accelerator = Accelerator(
+    mixed_precision="fp16"
+)
+
+controlnet = ControlNetModel.from_pretrained(
+    pretrained_model_name_or_path="nickpai/sdxl_light_caption_output",
+    subfolder="checkpoint-30000/controlnet",
+)
+
+pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+    pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0",
+    controlnet=controlnet
+)
+blip_processor = BlipProcessor.from_pretrained(
+    "Salesforce/blip-image-captioning-large",
+)
+blip_generator = BlipForConditionalGeneration.from_pretrained(
+    "Salesforce/blip-image-captioning-large",
+)
+
+pipe.to(accelerator.device)
+blip_generator.to(accelerator.device)
+
+def colorize_single_image(image, positive_prompt, negative_prompt, caption_generate):
+    image = PIL.Image.fromarray(image)
+
+    torch.cuda.empty_cache()
+    if caption_generate:
+        caption = blip_image_captioning(image=image, device=accelerator.device, processor=blip_processor, generator=blip_generator)
+    else:
+        caption = ""
+
+    original_size = image.size
+    control_image = image.convert("L").convert("RGB").resize((512, 512))
+    prompt = [positive_prompt + ", " + caption]
+    
+    colorized_image = pipe(prompt=prompt,
+                           num_inference_steps=5, 
+                           generator=torch.manual_seed(0),
+                           image=control_image,
+                           negative_prompt=negative_prompt).images[0]
+    result_image = apply_color(control_image, colorized_image)
+    result_image = result_image.resize(original_size)
+    return result_image, caption if caption_generate else gr.update(visible=False)
+
+# Hàm load models cho autoencoder và gan
+def load_autoencoder_model(auto_model_path):
     unet = UNetAuto(in_channels=1, out_channels=2).to(device)
     model = Autoencoder(unet).to(device)
     model.load_state_dict(torch.load(auto_model_path, map_location=device))
@@ -44,7 +104,7 @@ mobilenet_model = load_model(
     model_type='mobilenet'
 )
 
-unet_model = load_unet_model("weight/autoencoder.pt")
+autoencoder_model = load_autoencoder_model("weight/autoencoder.pt")
 
 # Transformations
 def preprocess_image(image):
@@ -67,68 +127,90 @@ def colorize_image(input_image, mode):
     with torch.no_grad():
         resnet_output = resnet_model.net_G(grayscale.unsqueeze(0))
         mobilenet_output = mobilenet_model.net_G(grayscale.unsqueeze(0))
-        unet_output = unet_model(grayscale.unsqueeze(0))
+        autoencoder_output = autoencoder_model(grayscale.unsqueeze(0))
     
     # Resize outputs to match the original size
     resnet_colorized = postprocess_image(grayscale, resnet_output, original_size)
     mobilenet_colorized = postprocess_image(grayscale, mobilenet_output, original_size)
-    unet_colorized = postprocess_image(grayscale, unet_output, original_size)
+    autoencoder_colorized = postprocess_image(grayscale, autoencoder_output, original_size)
     
     if mode == "ResNet":
         return resnet_colorized, None, None
     elif mode == "MobileNet":
         return None, mobilenet_colorized, None
     elif mode == "Unet":
-        return None, None, unet_colorized
+        return None, None, autoencoder_colorized
     elif mode == "Comparison":
-        return resnet_colorized, mobilenet_colorized, unet_colorized
+        return resnet_colorized, mobilenet_colorized, autoencoder_colorized
 
+def colorize_single_image(input_image, positive_prompt, negative_prompt, generate_caption):
+    
+    caption = "Generated Caption Example" if generate_caption else ""
+    return input_image, caption
 
-# Gradio Interface
 def gradio_interface():
-    with gr.Blocks() as demo:
-        # Input components
-        input_image = gr.Image(type="numpy", label="Upload an Image")
-        output_modes = gr.Radio(
-            choices=["ResNet", "MobileNet", "Unet", "Comparison"],
-            value="ResNet",
-            label="Output Mode"
-        )
-        
-        submit_button = gr.Button("Submit")
-
-        # Output components
-        with gr.Row():  # Place output images in a single row
-            resnet_output = gr.Image(label="Colorized Image (ResNet18)", visible=False)
-            mobilenet_output = gr.Image(label="Colorized Image (MobileNet)", visible=False)
-            unet_output = gr.Image(label="Colorized Image (Unet)", visible=False)
-
-        # Output mode logic
-        def update_visibility(mode):
-            if mode == "ResNet":
-                return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
-            elif mode == "MobileNet":
-                return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
-            elif mode == "Unet":
-                return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
-            elif mode == "Comparison":
-                return gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
-
-        # Dynamic event listener for output mode changes
-        output_modes.change(
-            fn=update_visibility,
-            inputs=[output_modes],
-            outputs=[resnet_output, mobilenet_output, unet_output]
-        )
-
-        # Submit logic
-        submit_button.click(
-            fn=colorize_image,
-            inputs=[input_image, output_modes],
-            outputs=[resnet_output, mobilenet_output, unet_output]
-        )
-
-    return demo
+    with gr.Blocks() as app:
+        with gr.Tab("Mode Colorization no Prompting"):
+            with gr.Blocks():
+                input_image = gr.Image(type="numpy", label="Upload an Image")
+                output_modes = gr.Radio(
+                    choices=["ResNet", "MobileNet", "Unet", "Comparison"],
+                    value="ResNet",
+                    label="Output Mode"
+                )
+                
+                submit_button = gr.Button("Submit")
+
+                with gr.Row():  # Place output images in a single row
+                    resnet_output = gr.Image(label="Colorized Image (ResNet18)", visible=False)
+                    mobilenet_output = gr.Image(label="Colorized Image (MobileNet)", visible=False)
+                    autoencoder_output = gr.Image(label="Colorized Image (Unet)", visible=False)
+
+                def update_visibility(mode):
+                    if mode == "ResNet":
+                        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
+                    elif mode == "MobileNet":
+                        return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
+                    elif mode == "Unet":
+                        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+                    elif mode == "Comparison":
+                        return gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
+
+                output_modes.change(
+                    fn=update_visibility,
+                    inputs=[output_modes],
+                    outputs=[resnet_output, mobilenet_output, autoencoder_output]
+                )
+
+                submit_button.click(
+                    fn=colorize_image,
+                    inputs=[input_image, output_modes],
+                    outputs=[resnet_output, mobilenet_output, autoencoder_output]
+                )
+
+        with gr.Tab("Stable Diffusion"):
+            with gr.Blocks():
+                sd_image = gr.Image(label="Upload a Color Image")
+                positive_prompt = gr.Textbox(label="Positive Prompt", placeholder="Text for positive prompt")
+                negative_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"
+                )
+                generate_caption = gr.Checkbox(label="Generate Caption?", value=False)
+                submit_sd = gr.Button("Generate")
+
+                sd_output_image = gr.Image(label="Colorized Image")
+                sd_caption = gr.Textbox(label="Captioning Result", show_copy_button=True, visible=False)
+
+                submit_sd.click(
+                    fn=colorize_single_image,
+                    inputs=[sd_image, positive_prompt, negative_prompt, generate_caption],
+                    outputs=[sd_output_image, sd_caption]
+                )
+
+    return app
+
+
 
 
 # Launch

stable.py

@@ -0,0 +1,115 @@
+# for image captioning
+import PIL
+import torch
+from torchvision import transforms
+import gradio as gr
+
+import transformers
+transformers.utils.move_cache()
+
+from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
+from transformers import BlipProcessor, BlipForConditionalGeneration
+from accelerate import Accelerator
+
+def remove_unlikely_words(prompt: str) -> str:
+    """
+    Removes unlikely words from a prompt.
+
+    Args:
+        prompt: The text prompt to be cleaned.
+
+    Returns:
+        The cleaned prompt with unlikely words removed.
+    """
+    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]
+
+    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",
+        "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",
+        "blurry photo,", "blurry,", "blurry photography,", "monochromatic photo",
+        "black - and - white photograph,", "black - and - white photograph", "black on white,",
+        "black on white", "black-and-white", "historical image,", "historical picture,", 
+        "historical photo,", "historical photograph,", "archival photo,", "taken in the early",
+        "taken in the late", "taken in the", "historic photograph,", "restored,", "restored", 
+        "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", 
+        "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)
+    
+    for word in unlikely_words:
+        prompt = prompt.replace(word, "")
+    return prompt
+
+def blip_image_captioning(image, device, processor, generator, conditional="a photography of"):
+    # Load the processor and model
+    if processor is None:
+        processor = BlipProcessor.from_pretrained(
+            "Salesforce/blip-image-captioning-large"
+        )
+    if generator is None:
+        model = BlipForConditionalGeneration.from_pretrained(
+            "Salesforce/blip-image-captioning-large",
+            torch_dtype=torch.float16
+        ).to(device)
+
+    # Prepare inputs
+    inputs = processor(
+        image, 
+        text=conditional, 
+        return_tensors="pt"
+    ).to(device)
+    
+    # Generate the caption
+    out = generator.generate(**inputs, max_new_tokens=20)  # Use max_new_tokens for better clarity
+    caption = processor.decode(out[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
+    caption = remove_unlikely_words(caption)
+
+    return caption
+
+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')
+    color_map_lab = color_map.convert('LAB')
+
+    # Split LAB channels
+    l, a , b = 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