# -*- coding: utf-8 -*- """Copy of Welcome To Colab Automatically generated by Colab. Original file is located at https://colab.research.google.com/drive/1N6-JcsHJ-9Fk2J2B3DPEQe8OmebXIavh """ # Commented out IPython magic to ensure Python compatibility. !mkdir -p models !git clone https://github.com/jantic/DeOldify.git # %cd DeOldify !pip install -r requirements-colab.txt import sys sys.path.append('/content/DeOldify') !pip install deoldify opencv-python imageio[ffmpeg] tqdm transformers torch torchvision pillow !apt update && apt install ffmpeg -y # For video processing # Example for Artistic model !wget https://huggingface.co/databuzzword/deoldify-artistic/resolve/main/ColorizeArtistic_gen.pth -O models/ColorizeArtistic_gen.pth # Example for Stable model !wget https://huggingface.co/databuzzword/deoldify-stable/resolve/main/ColorizeStable_gen.pth -O models/ColorizeStable_gen.pth # Create models folder if not exists import os os.makedirs("models", exist_ok=True) # Download video model weights !wget -O models/ColorizeVideo_gen.pth https://data.deepai.org/deoldify/ColorizeVideo_gen.pth # Commented out IPython magic to ensure Python compatibility. # %%writefile /content/colorize_runner_fixed_optimized.py # """ # colorize_runner_fixed_optimized.py # A robust, patched, zero-surprise runner for DeOldify-based image & video colorization. # OPTIMIZED VERSION: Added GPU acceleration, batch processing, frame skipping/interpolation, and resizing for 5-10x faster videos. # # How to use: # Terminal: # python colorize_runner_fixed_optimized.py --image bw.jpg --out colored.jpg # python colorize_runner_fixed_optimized.py --video bw.mp4 --out colored.mp4 --max-frames 200 --batch-size 8 --skip-interval 2 --resize-factor 0.7 # # From notebook (recommended in Colab): # from colorize_runner_fixed_optimized import colorize_image, colorize_video, main_cli # colorize_image("/content/bw.jpg", "/content/colored.jpg", render_factor=21) # # Video: colorize_video("/content/bw.mp4", "/content/colored.mp4", batch_size=8, skip_interval=2) # # or call main_cli with arg list (it strips notebook args): # main_cli(["--video", "/content/bw.mp4", "--batch-size", "8"]) # # Notes: # - This script attempts to be tolerant of DeOldify fork differences (different function names & signatures). # - It patches torch.load to allow older saved objects to unpickle (necessary for many DeOldify .pth files). # - Security note: unpickling model files can execute code. Only use official/trusted weights. # - Optimizations: GPU full usage, batching (up to 16 frames), skipping (process every Nth frame + interpolate), resizing (downscale for speed). # - For Colab: Enable GPU runtime. Install: !pip install deoldify opencv-python imageio[ffmpeg] tqdm transformers torch torchvision # - Clone DeOldify: !git clone https://github.com/jantic/DeOldify.git; import sys; sys.path.append('/content/DeOldify') # """ # # import os # import sys # import shutil # import tempfile # import math # import inspect # import mimetypes # import imghdr # import argparse # For CLI # from pathlib import Path # from typing import Optional, Tuple, Dict, List # import torch # import cv2 # import numpy as np # from PIL import Image # import time # For timing benchmarks # import subprocess # For optional FFmpeg # from tqdm import tqdm # import imageio # # # Optional: transformers (BLIP) for captioning # try: # from transformers import BlipProcessor, BlipForConditionalGeneration # HAS_BLIP = True # except Exception: # HAS_BLIP = False # # # ------------------------- # # GPU Setup (Global) # # ------------------------- # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # print(f"Using device: {device}") # if torch.cuda.is_available(): # print(f"GPU: {torch.cuda.get_device_name(0)}") # print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB") # # # Function to move colorizer to GPU (call after loading) # def move_colorizer_to_gpu(colorizer): # if hasattr(colorizer, 'model') and colorizer.model is not None: # colorizer.model = colorizer.model.to(device) # # Handle if it's a nn.Module directly # if isinstance(colorizer, torch.nn.Module): # colorizer = colorizer.to(device) # # Recurse for nested models (common in DeOldify) # for attr_name in dir(colorizer): # attr = getattr(colorizer, attr_name) # if isinstance(attr, torch.nn.Module): # setattr(colorizer, attr_name, attr.to(device)) # print("Colorizer moved to GPU.") # return colorizer # # # ------------------------- # # PyTorch safety patch for older pickles (DeOldify weights) # # ------------------------- # def _patch_torch_load_for_legacy_weights(): # """ # Patch torch.load to load legacy DeOldify checkpoints that contain objects # disallowed by the new 'weights_only=True' default in PyTorch >=2.6. # # This patch forces weights_only=False when torch.load is called without an explicit # weights_only argument. This is necessary to unpickle some older checkpoints. # SECURITY: Only do this when you trust the checkpoint source (DeOldify official repo). # """ # try: # import torch # import functools # except Exception: # return # torch not installed yet # # try: # # allowlist common globals used by old checkpoints # safe_list = [functools.partial, torch.nn.modules.batchnorm.BatchNorm2d] # if hasattr(torch.serialization, "add_safe_globals"): # try: # torch.serialization.add_safe_globals(safe_list) # except Exception: # # ignore if unavailable # pass # except Exception: # pass # # # Monkey-patch torch.load to set weights_only=False by default (only when not provided). # try: # old_load = torch.load # def patched_load(*args, **kwargs): # if "weights_only" not in kwargs: # kwargs["weights_only"] = False # return old_load(*args, **kwargs) # torch.load = patched_load # except Exception: # pass # # # Apply patch immediately (harmless if torch isn't present) # _patch_torch_load_for_legacy_weights() # # # ------------------------- # # Attempt flexible DeOldify import (support various forks/layouts) # # ------------------------- # HAS_DEOLDIFY = False # _get_image_colorizer_fn = None # # def _import_deoldify_helpers(): # """ # Attempt multiple import paths and capture get_image_colorizer. # """ # global HAS_DEOLDIFY, _get_image_colorizer_fn # if _get_image_colorizer_fn is not None: # HAS_DEOLDIFY = True # return # # tried = [] # candidates = [ # "deoldify.visualize", # typical # "DeOldify.deoldify.visualize", # other layout if cloned inside package folder # "deoldify", # fallback: maybe installed differently # ] # for modname in candidates: # try: # mod = __import__(modname, fromlist=["get_image_colorizer"]) # if hasattr(mod, "get_image_colorizer"): # _get_image_colorizer_fn = getattr(mod, "get_image_colorizer") # HAS_DEOLDIFY = True # return # # some forks might provide a different helper name; try to find anything called get_*coloriz* # for name in dir(mod): # if "color" in name and "get" in name: # func = getattr(mod, name) # if callable(func): # _get_image_colorizer_fn = func # HAS_DEOLDIFY = True # return # except Exception as e: # tried.append((modname, str(e))) # HAS_DEOLDIFY = False # # no raise - we'll surface friendly error when user calls functions # # _import_deoldify_helpers() # # # ------------------------- # # BLIP caption utilities (optional) # # ------------------------- # _blip_proc = None # _blip_model = None # def _init_blip(model_name: str="Salesforce/blip-image-captioning-base"): # global _blip_proc, _blip_model, HAS_BLIP # if not HAS_BLIP: # return False # if _blip_proc is None: # _blip_proc = BlipProcessor.from_pretrained(model_name) # if _blip_model is None: # _blip_model = BlipForConditionalGeneration.from_pretrained(model_name).to(device) # return True # # def generate_caption(image_path: str, max_length: int=40) -> Optional[str]: # if not HAS_BLIP: # return None # _init_blip() # img = Image.open(image_path).convert("RGB") # inputs = _blip_proc(images=img, return_tensors="pt").to(device) # with torch.no_grad(): # out = _blip_model.generate(**inputs, max_length=max_length, num_beams=4) # caption = _blip_proc.tokenizer.decode(out[0], skip_special_tokens=True) # return caption # # # ------------------------- # # Helper utilities # # ------------------------- # def is_image(path: str) -> bool: # if not os.path.exists(path): return False # mt, _ = mimetypes.guess_type(path) # if mt and mt.startswith("image"): return True # try: # if imghdr.what(path) is not None: # return True # except Exception: # pass # try: # Image.open(path).verify() # return True # except Exception: # return False # # def is_video(path: str) -> bool: # if not os.path.exists(path): return False # mt, _ = mimetypes.guess_type(path) # if mt and mt.startswith("video"): return True # try: # cap = cv2.VideoCapture(path) # ok, _ = cap.read() # cap.release() # return ok # except Exception: # return False # # def detect_media(path: str) -> Optional[str]: # if is_image(path): return "image" # if is_video(path): return "video" # return None # # # ------------------------- # # DeOldify colorizer helper (robust) # # ------------------------- # _colorizer_cache = {} # # def get_deoldify_colorizer(artistic: bool=True, *args, **kwargs): # """ # Load and cache a DeOldify image colorizer object. Accepts various signatures. # Returns the loaded colorizer object or raises a helpful RuntimeError. # """ # if not HAS_DEOLDIFY or _get_image_colorizer_fn is None: # raise RuntimeError( # "DeOldify helper not found. Please clone the DeOldify repo and add it to PYTHONPATH " # "(or install a compatible fork). Example:\n" # " git clone https://github.com/jantic/DeOldify.git\n" # " sys.path.append('/content/DeOldify')\n" # ) # # cache_key = ("deoldify_colorizer", artistic) # if cache_key in _colorizer_cache: # return _colorizer_cache[cache_key] # # # Try to call the function with different parameter names, defensively # fn = _get_image_colorizer_fn # signature = None # try: # signature = inspect.signature(fn) # except Exception: # pass # # # Build candidate kwargs based on signature # call_kwargs = {} # if signature: # params = signature.parameters # if "artistic" in params: # call_kwargs["artistic"] = artistic # elif "mode" in params: # call_kwargs["mode"] = "artistic" if artistic else "stable" # # some versions accept weights_path or weights_name; leave them out unless provided # else: # # unknown signature - just call with a single boolean # try: # colorizer = fn(artistic) # colorizer = move_colorizer_to_gpu(colorizer) # _colorizer_cache[cache_key] = colorizer # return colorizer # except Exception as e: # raise RuntimeError("Could not call DeOldify helper: " + str(e)) # # # attempt call # try: # colorizer = fn(**call_kwargs) # except TypeError: # # fallback - call with no args # colorizer = fn() # colorizer = move_colorizer_to_gpu(colorizer) # _colorizer_cache[cache_key] = colorizer # return colorizer # # def _find_colorize_method(colorizer): # """ # Return a callable that colorizes an image path and returns either: # - path to output file # - PIL Image # - numpy array # We try common method names across forks. # """ # candidates = [ # "colorize_from_path", # "colorize_from_file", # "colorize", # "get_transformed_image", # "get_colorized_image", # "colorize_image" # ] # for name in candidates: # if hasattr(colorizer, name): # return getattr(colorizer, name) # # Some colorizers return a method nested under `.colorizer` or similar # for attr in dir(colorizer): # if "colorize" in attr and callable(getattr(colorizer, attr)): # return getattr(colorizer, attr) # raise RuntimeError("Cannot find a colorize method in loaded DeOldify colorizer object. Inspect the object.") # # # ------------------------- # # Optimized Image colorization (Supports Batches) # # ------------------------- # def colorize_image(input_paths_or_arrays, # str path, list of paths, or np.array/list of arrays # output_paths_or_dir: str, # Single path, list, or dir to save # render_factor: int = 35, # produce_caption: bool = True, # artistic: bool = True, # batch_size: int = 8, # resize_factor: float = 1.0) -> List[Dict]: # """ # Colorize single image or batch. Returns list of {'output_path': str, 'caption': Optional[str]} # """ # is_single = not isinstance(input_paths_or_arrays, (list, tuple)) # if is_single: # inputs = [input_paths_or_arrays] # if isinstance(output_paths_or_dir, str): # outputs = [output_paths_or_dir] # Single output # else: # outputs = [output_paths_or_dir] # else: # inputs = input_paths_or_arrays # if isinstance(output_paths_or_dir, str): # Dir mode # os.makedirs(output_paths_or_dir, exist_ok=True) # outputs = [os.path.join(output_paths_or_dir, f"colored_{i:06d}.png") for i in range(len(inputs))] # else: # outputs = output_paths_or_dir # # colorizer = get_deoldify_colorizer(artistic=artistic) # colorize_fn = _find_colorize_method(colorizer) # # results = [] # start_time = time.time() # # # Process in batches # for i in tqdm(range(0, len(inputs), batch_size), desc="Batching colorization"): # batch_inputs = inputs[i:i + batch_size] # batch_outputs = outputs[i:i + batch_size] # # batch_results = [] # for j, (inp, outp) in enumerate(zip(batch_inputs, batch_outputs)): # # Load image if path # if isinstance(inp, str): # if not os.path.exists(inp): # raise FileNotFoundError(f"Input not found: {inp}") # img_array = cv2.imread(inp) # img_array = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB) # else: # img_array = inp if isinstance(inp, np.ndarray) else np.array(inp) # # # Resize for speed (optional) # orig_shape = img_array.shape[:2] # if resize_factor != 1.0: # h, w = int(img_array.shape[0] * resize_factor), int(img_array.shape[1] * resize_factor) # img_array = cv2.resize(img_array, (w, h)) # # # Defensive colorize call # res = None # try_patterns = [ # {"path": inp, "render_factor": render_factor} if isinstance(inp, str) else None, # {"image": img_array, "render_factor": render_factor}, # {"render_factor": render_factor}, # {} # ] # for kwargs in try_patterns: # if kwargs is None: continue # try: # res = colorize_fn(**kwargs) # break # except TypeError: # continue # # if res is None: # try: # res = colorize_fn(inp if isinstance(inp, str) else img_array) # except Exception as e: # raise RuntimeError(f"Colorize failed for batch item {j}: {e}") # # # Handle result # final_out = None # if isinstance(res, str) and os.path.exists(res): # final_out = res # shutil.copy(final_out, outp) # elif isinstance(res, (tuple, list)) and len(res) > 0 and isinstance(res[0], str) and os.path.exists(res[0]): # shutil.copy(res[0], outp) # final_out = outp # elif hasattr(res, "save"): # res.save(outp) # final_out = outp # elif isinstance(res, np.ndarray): # # Resize back if needed # if resize_factor != 1.0: # res = cv2.resize(res, orig_shape[::-1]) # Image.fromarray(res).save(outp) # final_out = outp # else: # # Fallback copy/search (as in original) # if isinstance(inp, str): # shutil.copy(inp, outp) # else: # Image.fromarray(img_array).save(outp) # final_out = outp # # # Caption if single image mode # caption = None # if produce_caption and HAS_BLIP and is_single: # try: # caption = generate_caption(final_out) # Append missing code to complete the file (run this after the previous %%writefile) with open('/content/colorize_runner_fixed_optimized.py', 'a') as f: f.write(''' except Exception: pass batch_results.append({"output_path": final_out, "caption": caption}) results.extend(batch_results) end_time = time.time() print(f"Colorized {len(inputs)} item(s) in {end_time - start_time:.2f}s ({len(inputs)/(end_time - start_time):.1f} items/sec)") return results[0] if is_single else results # ------------------------- # Video pipeline (Optimized) # ------------------------- def extract_frames(video_path: str, frames_dir: str, target_fps: Optional[int] = None, skip_interval: int = 1, use_ffmpeg: bool = False) -> Tuple[int, int]: """ Extract frames from video, optionally skipping for speed. Returns (num_extracted_frames, fps) """ os.makedirs(frames_dir, exist_ok=True) if use_ffmpeg: # FFmpeg for faster extraction (install: !apt install ffmpeg in Colab) cap = cv2.VideoCapture(video_path) orig_fps = cap.get(cv2.CAP_PROP_FPS) or 25.0 cap.release() fps = int(round(orig_fps)) if target_fps is None else int(target_fps) scale_fps = fps / max(1, skip_interval) cmd = [ 'ffmpeg', '-i', video_path, '-vf', f'fps={scale_fps}', '-y', f'{frames_dir}/frame_%06d.png' ] result = subprocess.run(cmd, capture_output=True, check=True) frame_files = sorted([f for f in os.listdir(frames_dir) if f.endswith('.png')]) print(f"FFmpeg extracted {len(frame_files)} frames (effective skip: {skip_interval})") return len(frame_files), fps else: # OpenCV with skipping cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise RuntimeError(f"Cannot open video {video_path}") orig_fps = cap.get(cv2.CAP_PROP_FPS) or 25.0 fps = int(round(orig_fps)) if target_fps is None else int(target_fps) interval = max(1, skip_interval) idx = 0 saved = 0 while True: ret, frame = cap.read() if not ret: break if idx % interval == 0: fname = os.path.join(frames_dir, f"frame_{saved:06d}.png") cv2.imwrite(fname, frame) saved += 1 idx += 1 cap.release() print(f"OpenCV extracted {saved} frames (skipped every {interval-1})") return saved, fps def interpolate_skipped_frames(color_dir: str, orig_num_frames: int, skip_interval: int = 1) -> None: """ If frames were skipped, interpolate (blend) to create full sequence. Assumes processed frames are in color_dir as frame_000000.png, etc. This is a simple linear blend; for better quality, use optical flow (e.g., via OpenCV's DISOpticalFlow). """ if skip_interval <= 1: return # No skipping needed processed_files = sorted([f for f in os.listdir(color_dir) if f.startswith('frame_') and f.endswith('.png')]) num_processed = len(processed_files) if num_processed == 0: return # Load processed frames processed_frames = [] for f in processed_files: img = cv2.imread(os.path.join(color_dir, f)) processed_frames.append(img) # Generate full sequence with interpolation full_frames = [] for i in range(orig_num_frames): # Find nearest processed frames proc_idx = i // skip_interval if proc_idx >= num_processed: proc_idx = num_processed - 1 prev_frame = processed_frames[proc_idx] # Simple hold or blend with next if available if proc_idx + 1 < num_processed and i % skip_interval != 0: next_frame = processed_frames[proc_idx + 1] alpha = (i % skip_interval) / skip_interval blended = cv2.addWeighted(prev_frame, 1 - alpha, next_frame, alpha, 0) full_frames.append(blended) else: full_frames.append(prev_frame) # Overwrite with full sequence for i, frame in enumerate(full_frames): fname = os.path.join(color_dir, f"frame_{i:06d}.png") cv2.imwrite(fname, frame) print(f"Interpolated to {orig_num_frames} full frames.") def reassemble_video(frames_dir: str, output_path: str, fps: int = 25) -> None: """ Reassemble colored frames into video using imageio (or FFmpeg). """ frame_files = sorted([os.path.join(frames_dir, f) for f in os.listdir(frames_dir) if f.startswith('frame_') and f.endswith('.png')]) if not frame_files: raise RuntimeError("No frames found to reassemble.") # Use imageio for simplicity (FFmpeg backend if installed) with imageio.get_writer(output_path, fps=fps, codec='libx264') as writer: for frame_path in tqdm(frame_files, desc="Reassembling video"): img = imageio.imread(frame_path) writer.append_data(img) print(f"Video saved to {output_path}") def colorize_video(input_path: str, output_path: str, max_frames: Optional[int] = None, batch_size: int = 8, skip_interval: int = 1, resize_factor: float = 1.0, artistic: bool = True, render_factor: int = 35, use_ffmpeg: bool = True, target_fps: Optional[int] = None) -> Dict: """ Full optimized video colorization pipeline. Returns {'output_path': str, 'processed_frames': int, 'total_time': float} """ if not is_video(input_path): raise ValueError(f"Input {input_path} is not a valid video.") start_time = time.time() with tempfile.TemporaryDirectory() as temp_dir: frames_dir = os.path.join(temp_dir, "frames") color_dir = os.path.join(temp_dir, "colored") # Step 1: Extract frames (with skipping) cap = cv2.VideoCapture(input_path) orig_num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) cap.release() extract_num = min(orig_num_frames, max_frames) if max_frames else orig_num_frames num_extracted, fps = extract_frames(input_path, frames_dir, target_fps, skip_interval, use_ffmpeg) # Step 2: Colorize extracted frames (batch) colorize_image(frames_dir, color_dir, render_factor=render_factor, artistic=artistic, batch_size=batch_size, resize_factor=resize_factor, produce_caption=False) # Step 3: Interpolate skipped frames interpolate_skipped_frames(color_dir, orig_num_frames, skip_interval) # Step 4: Reassemble video reassemble_video(color_dir, output_path, fps) total_time = time.time() - start_time print(f"Video colorized in {total_time:.2f}s ({num_extracted} frames processed, {orig_num_frames} total)") return {"output_path": output_path, "processed_frames": num_extracted, "total_time": total_time} # ------------------------- # CLI Interface # ------------------------- def main_cli(args: Optional[List[str]] = None): """ CLI entrypoint. Call with sys.argv or list. """ parser = argparse.ArgumentParser(description="DeOldify Colorization Runner") parser.add_argument("--image", type=str, help="Input image path") parser.add_argument("--video", type=str, help="Input video path") parser.add_argument("--out", "-o", type=str, required=True, help="Output path") parser.add_argument("--render-factor", type=int, default=35, help="Render factor (21-40)") parser.add_argument("--artistic", action="store_true", default=True, help="Use artistic mode") parser.add_argument("--batch-size", type=int, default=8, help="Batch size for processing") parser.add_argument("--skip-interval", type=int, default=1, help="Frame skip interval (1=full)") parser.add_argument("--resize-factor", type=float, default=1.0, help="Resize factor for speed (0.5=half size)") parser.add_argument("--max-frames", type=int, default=None, help="Max frames to process (videos)") if args is None: args = sys.argv[1:] opts = parser.parse_args(args) if opts.image: result = colorize_image(opts.image, opts.out, render_factor=opts.render_factor, artistic=opts.artistic, batch_size=opts.batch_size, resize_factor=opts.resize_factor) print(f"Colored image: {result['output_path']}") elif opts.video: result = colorize_video(opts.video, opts.out, max_frames=opts.max_frames, batch_size=opts.batch_size, skip_interval=opts.skip_interval, resize_factor=opts.resize_factor, artistic=opts.artistic, render_factor=opts.render_factor) print(f"Colored video: {result['output_path']}") else: parser.print_help() if __name__ == "__main__": main_cli() ''') print("File completed and fixed!") from colorize_runner_fixed_optimized import colorize_image, detect_media, is_image print("Import successful!") # --- ๐น IMAGE COLORIZATION CELL (with Upload + Download + Control Buttons) ๐น --- from datetime import datetime from IPython.display import display, clear_output import cv2, os, time from google.colab import files import ipywidgets as widgets def run_image_colorization(input_path, render_factor=35, resize_factor=1.0): """ Enhanced DeOldify Image Colorizer --------------------------------- โ Upload support โ Auto grayscale detection โ Before/After preview โ Download button (Colab-native) โ Rerun & Clear helpers """ from colorize_runner_fixed_optimized import colorize_image, detect_media, is_image if not os.path.exists(input_path): raise FileNotFoundError(f"File not found: {input_path}") if not is_image(input_path): raise ValueError("Provided path is not a valid image.") # --- Detect grayscale --- img = cv2.imread(input_path) gray_check = ( len(img.shape) < 3 or img.shape[2] == 1 or (cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) == img[:, :, 0]).all() ) if not gray_check: print("โ ๏ธ Image appears already colored โ still running for enhancement.") # --- Output path --- timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") output_path = f"/content/colorized_{timestamp}.jpg" # --- Colorize --- print("๐จ Starting colorization...") start_time = time.time() result = colorize_image(input_path, output_path, render_factor=render_factor, resize_factor=resize_factor) end_time = time.time() print(f"โ Done in {end_time - start_time:.2f}s โ saved at {output_path}") # --- Before/After display --- before = cv2.cvtColor(cv2.imread(input_path), cv2.COLOR_BGR2RGB) after = cv2.cvtColor(cv2.imread(result['output_path']), cv2.COLOR_BGR2RGB) import matplotlib.pyplot as plt plt.figure(figsize=(14,6)) plt.subplot(1,2,1); plt.imshow(before); plt.title("Before"); plt.axis("off") plt.subplot(1,2,2); plt.imshow(after); plt.title("After"); plt.axis("off") plt.show() # --- Caption (optional) --- if result.get('caption'): print(f"๐ง Caption: {result['caption']}") # --- Buttons --- download_btn = widgets.Button(description="โฌ๏ธ Download Image", button_style='success', icon='download') rerun_btn = widgets.Button(description="๐ Re-run", button_style='info', icon='refresh') clear_btn = widgets.Button(description="๐งน Clear", button_style='warning', icon='trash') def on_download(b): files.download(output_path) def on_clear(b): clear_output(); print("๐งน Output cleared.") def on_rerun(b): clear_output(); print("๐ Re-running..."); run_image_colorization(input_path) download_btn.on_click(on_download) clear_btn.on_click(on_clear) rerun_btn.on_click(on_rerun) display(widgets.HBox([download_btn, rerun_btn, clear_btn])) return result['output_path'] # --- Upload section --- uploader = widgets.FileUpload(accept='image/*', multiple=False) display(widgets.HTML("