Upload app.py

app.py

@@ -9,8 +9,8 @@ import sys
 from PIL import Image, ImageFilter, ImageDraw
 from diffusers import StableDiffusionInpaintPipeline
 from transformers import BlipProcessor, BlipForConditionalGeneration
-from diffusers import StableDiffusionImg2ImgPipeline
-
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
 import torchvision.transforms.functional as F
 sys.modules["torchvision.transforms.functional_tensor"] = F
 
@@ -53,18 +53,21 @@ print("Real-ESRGAN ready.")
 print("Loading SD Inpainting...")
 
 inpaint = StableDiffusionInpaintPipeline.from_pretrained(
-    "runwayml/stable-diffusion-inpainting",
+    "stabilityai/stable-diffusion-2-inpainting",
     torch_dtype=torch.float16,
+    variant= "fp16"
 )
 print("SD Inpainting ready.")
 
-# COLORING- IMG2IMG Pipeline
-colour= StableDiffusionImg2ImgPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    safety_checker=None,
-    requires_safety_checker=False,
+# COLORING
+colorizer = pipeline(
+    Tasks.image_colorization,
+    model="damo/cv_ddcolor_image-colorization"
+    # DDColor = dedicated colorization model
+    # Trained specifically to add color without changing structure
 )
-print("colour model ready!")
+
+print("Colorizer ready.")
 
 
 #  BLIP 
@@ -103,7 +106,7 @@ def get_caption(image):
     caption = blip_processor.decode(output[0], skip_special_tokens=True)
     return caption
 
-def build_prompt(caption, image):
+def build_prompt(caption, image): #for outpaint
     bw = is_greyscale(image)
     style_hint = (
         "black and white photography, monochrome, greyscale, "
@@ -117,9 +120,105 @@ def build_prompt(caption, image):
         f"same atmosphere, high quality"
     )
 
+
+
+@spaces.GPU
+def enhance_image(image, scale_factor):
+    if image is None:
+        raise gr.Error("Please upload an image first.")
+
+    # Move models to GPU — happens inside the decorated function
+    # because GPU is only available here
+    enhancer.device = torch.device("cuda")
+    enhancer.half   = True
+
+    image_array = np.array(image)
+    image_array = image_array[:, :, :3]
+    outscale    = 4 if scale_factor == "4x" else 2
+
+    try:
+        output_array, _ = enhancer.enhance(image_array, outscale=outscale)
+    except RuntimeError as e:
+        raise gr.Error(f"Enhancement failed: {e}. Try a smaller image.")
+
+    output_rgb   = output_array[:, :, ::-1]
+    output_image = Image.fromarray(output_rgb)
+
+    original_size = f"{image.width}×{image.height}"
+    new_size      = f"{output_image.width}×{output_image.height}"
+
+    return output_image, f"Original: {original_size} → Enhanced: {new_size}"
+
+@spaces.GPU
+def colour_image(image, strength):
+    
+    if Image is None:
+        raise gr.Error("Please Upload the Image")
+
+    # Convert PIL to numpy RGB
+    img_array = image.convert("RGB")
+    # Resize so longest side = 512
+    # DDColor works best at this resolution
+    target = 512
+    ratio  = min(target / image.width, target / image.height)
+    w      = int(image.width  * ratio)
+    h      = int(image.height * ratio)
+    image  = image.resize((w, h), Image.LANCZOS)
+    # LANCZOS = high quality resampling, preserves sharp edges
+    # DDColor expects BGR (OpenCV format)
+    # Convert PIL → numpy BGR for DDColor
+    img_rgb = np.array(image)
+    # np.array(PIL) = numpy array shape (h, w, 3) in RGB
+    img_bgr = img_rgb[:, :, ::-1]
+
+    #run colorizer
+    result= colorizer(img_bgr)
+    output_bgr= result["output_img"]
+
+    # Convert BGR → RGB for PIL
+    output_rgb = output_bgr[:, :, ::-1]
+
+    # Apply strength blending
+    if strength < 1.0:
+        grey    = np.array(image.convert("L"))
+        # convert("L") = convert to greyscale (single channel)
+        grey_3ch = np.stack([grey, grey, grey], axis=-1)
+        # stack same channel 3 times → (h, w, 3) array
+        # needed to match output_rgb shape for blending
+        output_rgb = (
+            strength * output_rgb.astype(float) +
+            (1 - strength) * grey_3ch.astype(float)
+        ).astype(np.uint8)
+        # Linear blend: weighted average of colored and grey
+        # .astype(np.uint8) = convert back to 0-255 integers
+        return Image.fromarray(output_rgb)
+
+
+def get_caption(image):
+    inputs  = blip_processor(image.convert("RGB"), return_tensors="pt").to("cuda", torch.float16)
+    output  = blip_model.generate(**inputs, max_new_tokens=60)
+    caption = blip_processor.decode(output[0], skip_special_tokens=True)
+    return caption
+
+
+def build_prompt(caption):
+    return (
+        f"seamless continuation of a scene with {caption}, "
+        f"same exact lighting, same exact style, "
+        f"same background, extending the existing scene naturally, "
+        f"no new objects, no new subjects, only environment continuation"
+    )
+
+
 def extend_one_side(image, direction, pixels, prompt, negative_prompt):
-    orig_w = image.width
-    orig_h = image.height
+    # CORE LOGIC — extends one side by a SMALL number of pixels
+    # Called multiple times in small steps instead of one big jump
+    # Small steps = SD always has lots of context = coherent output
+    
+    from PIL import ImageDraw
+
+    orig_w  = image.width
+    orig_h  = image.height
 
     new_w   = orig_w
     new_h   = orig_h
@@ -129,19 +228,24 @@ def extend_one_side(image, direction, pixels, prompt, negative_prompt):
     if direction == "left":
         new_w   = orig_w + pixels
         paste_x = pixels
+
     elif direction == "right":
         new_w   = orig_w + pixels
         paste_x = 0
+
     elif direction == "top":
         new_h   = orig_h + pixels
         paste_y = pixels
+
     elif direction == "bottom":
         new_h   = orig_h + pixels
         paste_y = 0
 
+    # Round to multiple of 8 — SD requirement
     new_w = (new_w // 8) * 8
     new_h = (new_h // 8) * 8
 
+    # Recalculate after rounding
     if direction in ["left", "right"]:
         pixels = new_w - orig_w
     else:
@@ -152,27 +256,54 @@ def extend_one_side(image, direction, pixels, prompt, negative_prompt):
     if direction == "top":
         paste_y = pixels
 
+    # Build black canvas and paste original into it
     canvas = Image.new("RGB", (new_w, new_h), (0, 0, 0))
     canvas.paste(image, (paste_x, paste_y))
 
+    # ── Build mask ────────────────────────────────────────────
     mask = Image.new("L", (new_w, new_h), 255)
+    # All white = generate everywhere
     draw = ImageDraw.Draw(mask)
+
     feather = 30
+    # This ensures the black region doesn't touch the very edge of original
+    # leaving a thin white strip that GaussianBlur will turn into a gradient
+
 
     if direction == "left":
-        draw.rectangle([paste_x + feather, feather, new_w - feather, new_h - feather], fill=0)
+        draw.rectangle([
+            paste_x + feather,  feather,
+            new_w   - feather,  new_h - feather
+        ], fill=0)
+
     elif direction == "right":
-        draw.rectangle([feather, feather, orig_w - feather, new_h - feather], fill=0)
-    elif direction == "top":
-        draw.rectangle([feather, paste_y + feather, new_w - feather, new_h - feather], fill=0)
-    elif direction == "bottom":
-        draw.rectangle([feather, feather, new_w - feather, orig_h - feather], fill=0)
+        draw.rectangle([
+            feather, feather,
+            orig_w  - feather,  new_h - feather
+        ], fill=0)
 
-    mask = mask.filter(ImageFilter.GaussianBlur(radius=30))
+    elif direction == "top":
+        draw.rectangle([
+            feather,            paste_y + feather,
+            new_w   - feather,  new_h   - feather
+        ], fill=0)
 
-    sd_size   = 512
+    elif direction == "bottom":
+        draw.rectangle([
+            feather,            feather,
+            new_w   - feather,  orig_h  - feather
+        ], fill=0)
+
+    
+    #Apply GaussianBlur to the mask AFTER drawing
+    # This is what creates REAL feathering — a soft gradient at the boundary
+    mask= mask.filter(ImageFilter.GaussianBlur(radius=30))
+
+    
+    sd_size   = 768
     canvas_sd = canvas.resize((sd_size, sd_size), Image.LANCZOS)
     mask_sd   = mask.resize((sd_size, sd_size), Image.LANCZOS)
+    # LANCZOS = high quality downsampling, preserves sharp details
 
     result = inpaint(
         prompt              = prompt,
@@ -185,132 +316,83 @@ def extend_one_side(image, direction, pixels, prompt, negative_prompt):
         negative_prompt     = negative_prompt,
     )
 
-    generated_512  = result.images[0]
+    generated_512 = result.images[0]
+
     generated_full = generated_512.resize((new_w, new_h), Image.LANCZOS)
-    # No hard paste — GaussianBlur mask handles the boundary softly
     return generated_full
 
-@spaces.GPU
-def enhance_image(image, scale_factor):
-    if image is None:
-        raise gr.Error("Please upload an image first.")
-
-    # Move models to GPU — happens inside the decorated function
-    # because GPU is only available here
-    enhancer.device = torch.device("cuda")
-    enhancer.half   = True
-
-    image_array = np.array(image)
-    image_array = image_array[:, :, :3]
-    outscale    = 4 if scale_factor == "4x" else 2
-
-    try:
-        output_array, _ = enhancer.enhance(image_array, outscale=outscale)
-    except RuntimeError as e:
-        raise gr.Error(f"Enhancement failed: {e}. Try a smaller image.")
-
-    output_rgb   = output_array[:, :, ::-1]
-    output_image = Image.fromarray(output_rgb)
-
-    original_size = f"{image.width}×{image.height}"
-    new_size      = f"{output_image.width}×{output_image.height}"
-
-    return output_image, f"Original: {original_size} → Enhanced: {new_size}"
-
-@spaces.GPU
-def colour_image(image, prompt, strength):
-    colour.to("cuda", torch.float16)
-    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= float(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:
         raise gr.Error("Please upload an image first.")
-
-    # Move models to GPU inside the decorated function
     inpaint.to("cuda")
     blip_model.to("cuda")
 
-    target_side = 512
-    ratio       = min(target_side / image.width, target_side / image.height)
-    new_size    = (int(image.width * ratio), int(image.height * ratio))
-    image       = image.resize(new_size, Image.LANCZOS)
+    # Resize input to max 512 on longest side
+    max_side = 512
+    ratio    = min(max_side / image.width, max_side / image.height)
+    
+    new_size = (int(image.width * ratio), int(image.height * ratio))
+    image    = image.resize(new_size, Image.LANCZOS)
 
     progress(0.05, desc="Analyzing image with BLIP...")
 
-    blip_caption = custom_prompt.strip() if custom_prompt.strip() else get_caption(image)
-    prompt       = build_prompt(blip_caption, image)
+    if custom_prompt.strip():
+        blip_caption = custom_prompt.strip()
+    else:
+        blip_caption = get_caption(image)
+
+    prompt = build_prompt(blip_caption)
 
     negative_prompt = (
         "blurry, bad quality, watermark, text, "
-        "new person, new face, extra people, "
+        "new person, new face, new subject, extra people, "
         "colorful, vibrant colors, color photography, "
-        "duplicate, border, frame, seam, visible edge, "
-        "distorted, inconsistent style"
+        "duplicate, tiled, repeated pattern, border, frame, "
+        "seam, visible edge, abrupt change, inconsistent, "
+        "distorted, unnatural, different style, different era"
     )
 
+    # ── KEY CHANGE: small fixed step size, multiple passes ────
     STEP_PX = 64
-    extend  = extend_percent / 100.0
-    h_total = int(image.width  * extend)
-    v_total = int(image.height * extend)
+    # Add only 64 pixels per pass instead of the full extension at once
+    # 64px out of a 512px image = 12.5% of canvas = SD has 87.5% context
+    # Compare to before: 35% extension = SD only had 65% context
+    # More context = SD understands the scene = coherent generation
+
+    # Calculate total pixels needed per direction
+    extend   = extend_percent / 100.0
+    h_total  = int(image.width  * extend)
+    # Total horizontal pixels to add on each side
+    v_total  = int(image.height * extend)
+    # Total vertical pixels to add on each side
 
     def make_passes(side, total_px):
-        passes    = []
+        # Break total_px into multiple STEP_PX passes
+        # Example: total=180px, STEP_PX=64 → passes of [64, 64, 52]
+        passes = []
         remaining = total_px
         while remaining > 0:
             step = min(STEP_PX, remaining)
+            # min() = don't overshoot — last step may be smaller
             passes.append((side, step))
             remaining -= step
         return passes
+        # Returns list of (direction, pixels) tuples
+        # Each tuple = one call to extend_one_side
 
+    # Build full pass list based on chosen direction
     if direction == "Horizontal":
         passes = make_passes("right", h_total) + make_passes("left", h_total)
+        # Extend right in small steps, then left in small steps
+
     elif direction == "Vertical":
         passes = make_passes("bottom", v_total) + make_passes("top", v_total)
+        # Bottom first (more grounded content), then top
+
     else:
+        # Both — all four sides in small steps
         passes = (
             make_passes("bottom", v_total) +
             make_passes("top",    v_total) +
@@ -322,20 +404,24 @@ def outpaint_image(image, direction, extend_percent, custom_prompt, progress=gr.
     current_image = image
 
     for i, (side, px) in enumerate(passes):
-        progress(
-            0.1 + 0.85 * (i / total_passes),
-            desc=f"Pass {i+1}/{total_passes} — extending {side} by {px}px"
-        )
+        progress_val = 0.1 + 0.85 * (i / total_passes)
+        progress(progress_val, desc=f"Pass {i+1}/{total_passes} — extending {side} by {px}px")
+
         current_image = extend_one_side(
-            current_image, side, px, prompt, negative_prompt
+            current_image,
+            side,
+            px,
+            prompt,
+            negative_prompt
         )
+        # Each pass returns a slightly larger image
+        # Next pass uses that as input — builds on previous result
 
     progress(1.0, desc="Done!")
 
-    bw_note = " [B&W detected]" if is_greyscale(image) else ""
     return (
         current_image,
-        f"Caption{bw_note}:\n{blip_caption}\n\nPrompt:\n{prompt}"
+        f"BLIP caption:\n{blip_caption}\n\nFull prompt:\n{prompt}"
     )
 
 # GRADIO UI
@@ -379,9 +465,6 @@ with gr.Blocks(title="CanvasAI — Enhance & Outpaint") as demo:
                         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,
@@ -396,11 +479,10 @@ with gr.Blocks(title="CanvasAI — Enhance & Outpaint") as demo:
                     )
                 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]
+                inputs= [col_input, col_strength],
+                outputs= [col_output]
             )
 
         with gr.Tab(" Outpaint"):