Upload app.py

app.py

@@ -9,6 +9,7 @@ import sys
 from PIL import Image, ImageFilter, ImageDraw
 from diffusers import StableDiffusionInpaintPipeline
 from transformers import BlipProcessor, BlipForConditionalGeneration
+from diffusers import StableDiffusionImg2ImgPipeline
 
 import torchvision.transforms.functional as F
 sys.modules["torchvision.transforms.functional_tensor"] = F
@@ -57,6 +58,18 @@ inpaint = StableDiffusionInpaintPipeline.from_pretrained(
 )
 print("SD Inpainting ready.")
 
+# COLORING- IMG2IMG Pipeline
+colour= StableDiffusionImg2ImgPipeline(
+    vae= inpaint.vae,
+    text_encoder= inpaint.text_encoder,
+    tokenizer= inpaint.tokenizer,
+    unet= inpaint.unet,
+    scheduler= inpaint.scheduler,
+    safety_checker= None,
+    feature_extractor= None,
+    requires_safety_checker= False,
+).to("cuda")
+
 
 #  BLIP 
 print("Loading BLIP...")
@@ -207,6 +220,53 @@ def enhance_image(image, scale_factor):
 
     return output_image, f"Original: {original_size} → Enhanced: {new_size}"
 
+@spaces.GPU
+def colour_image(image, prompt, strength):
+    if Image is None:
+        raise gr.Error("Please Upload the Image")
+    # Convert to RGB
+    image= image.convert("RGB")
+    # Resize to 512 for SD
+    target= 512
+    ratio= min(target/ image.width, target/image.height)
+    image= image.resize(
+        (int(image.width * ratio), int(image.height * ratio)),
+        Image.LANCZOS
+    )
+
+    #Build Prompt
+    if not prompt.strip():
+        # auto generate with BLIP
+        caption= get_caption(image)
+        full_prompt= (
+            f"colorized photograph, {caption}, "
+            f"natural realistic colors, vivid, sharp, "
+            f"professional color grading, high quality"
+        )
+    else:
+        full_prompt=  (
+            f"colorized photograph, {prompt}, "
+            f"natural realistic colors, high quality"
+        )
+
+    negative_prompt= (
+        "black and white, greyscale, monochrome, "
+        "blurry, bad quality, oversaturated, unnatural colors"
+    )
+
+    output= colour(
+        prompt= full_prompt,
+        image= image,
+        strength= strength,
+        negative_prompt= negative_prompt,
+        num_inference_steps= 30,
+        guidance_scale= 7.5
+    )
+
+    result= output.images[0]
+    return result, f"Prompt used: \n {full_prompt}"
+
+
 @spaces.GPU
 def outpaint_image(image, direction, extend_percent, custom_prompt, progress=gr.Progress()):
     if image is None:
@@ -312,6 +372,39 @@ with gr.Blocks(title="CanvasAI — Enhance & Outpaint") as demo:
                 outputs=[enh_output, enh_info]
             )
 
+        with gr.Tab("Colourize"):
+            gr.Markdown("Upload a black and white image. "
+        "AI adds natural, realistic colours while preserving the original structure.")
+            with gr.Row:
+                with gr.Column:
+                    col_input= gr.Image(
+                        label="Upload Image",
+                        type="PIL",
+                    )
+                    col_prompt= gr.Textbox(label= "Custom Prompt(optional)", 
+                                           placeholder="Leave empty for auto-detection, or type: portrait of a young woman in 1960s clothing",
+                                           lines=2)
+                    col_strength= gr.Slider(
+                        minimum= 0.3,
+                        maximum= 0.7,
+                        value= 0.5,
+                        step= 0.05,
+                        label= "Colorization Strength",
+                        info= "Low = subtle tint, stays close to original. High = vivid colours."
+                    )
+
+                    col_btn= gr.Button(
+                        "Colorize", variant="primary"
+                    )
+                with gr.Column():
+                    col_output= gr.Image(label="Colorized Result", type="pil", interactive=False)
+                    col_caption = gr.Textbox(label="Prompt Used", interactive=False, lines=3)
+            col_btn.click(
+                fn= colour_image,
+                inputs= [col_input, col_prompt, col_strength],
+                outputs= [col_output, col_caption]
+            )
+
         with gr.Tab(" Outpaint"):
             gr.Markdown(
                 "Upload an image and extend it in any direction. "