Using Stable Diffision models for Colorization

[didris]

installation was not complicated, but I also could not get it to work this way.
I installed it on embedded python - there are quite a few models - C:\Users\YOUR_USERNAME\.cache\huggingface\hub\ - 30.8 GB

Then I created two python files in the installation folder E:\DiTServerRPC
rpc_wrapper.py

from dit_client_example import main as _run_single def colorize_image(image_path: str):     """     Wrapper around working CLI logic.     Returns bytes of output image.     """     # It just uses the existing working flow     # (we do not touch the RPC logic)     return _run_single(image_path)

batch_colorize_turbo.py

import xmlrpc.client import io from pathlib import Path from PIL import Image import time # ========================= # CONFIG # ========================= HOST = "127.0.0.1" PORT = 8765 INPUT_DIR = Path("assets") OUTPUT_DIR = Path("output") OUTPUT_DIR.mkdir(exist_ok=True) SUPPORTED = {".png", ".jpg", ".jpeg", ".webp"} PROMPT = "Colorize this photo, natural skin tones, cinematic lighting" STEPS = 4 # ========================= # HELPERS # ========================= def pil_to_bytes(img):     buf = io.BytesIO()     img.save(buf, format="PNG")     return buf.getvalue() def bytes_to_pil(data):     raw = data.data if hasattr(data, "data") else data     return Image.open(io.BytesIO(raw)).convert("RGB") # ========================= # MAIN # ========================= def main():     print("🚀 Connecting...")     server = xmlrpc.client.ServerProxy(         f"http://{HOST}:{PORT}/",         use_builtin_types=True     )     server.ping()     print("✅ Server OK")     images = sorted([         p for p in INPUT_DIR.iterdir()         if p.suffix.lower() in SUPPORTED     ])     print(f"🚀 Found {len(images)} images")     start = time.perf_counter()     # 🔥 IMPORTANT: SERIAL (GPU-safe)     for img_path in images:         print(f"[RPC] {img_path.name}")         img = Image.open(img_path).convert("RGB")         try:             result = server.colorize_frame(                 pil_to_bytes(img),                 PROMPT,                 0,                 STEPS             )             if not result["ok"]:                 print(f"❌ ERROR: {result['msg']}")                 continue             out = bytes_to_pil(result["data"])             out_path = OUTPUT_DIR / img_path.name             out.save(out_path)             print(f"✅ Saved: {out_path}")         except Exception as e:             print(f"❌ ERROR {img_path.name}: {e}")     print(f"\n⚡ TOTAL TIME: {time.perf_counter() - start:.2f}s") if __name__ == "__main__":     main()

I added this function to the existing dit_client_example.py

colorize_image()

I run it with this command in powershell

cd E:\DiTServerRPC E:\python_embeded\python.exe dit_rpc_server.py --load-pipeline --pipeline-config qwen_config_fp4.json

then in another powershell window

cd E:\DiTServerRPC E:\python_embeded\python.exe batch_colorize_turbo.py

What it does is the following - it takes the frames one by one from the folder "E:\DiTServerRPC\assets" and processes them one after the other automatically in the folder "E:\DiTServerRPC\output" while keeping the same name, resolution and jpg format.

The result is an average of 9 seconds per image, uses an average of 23 GB of gpu memory and 39 GB of ram during the process /my card is rtx5090/, It is probably possible to improve the time, but it will be at the expense of quality. It does not colorize quite evenly - on different frames the same thing sometimes colors it differently - probably a lot depends on what is set in the prompt.

Once again, thanks to Dan and Selur for what they have done.