From 008c30bca834447cad4f0aed427178cd155d4703 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Sun, 24 May 2026 12:33:00 +0300 Subject: chore: remove photo/ComfyUI top-level files Delete hyperstack-vm-photo.toml, photo-enhance.rb, photo-enhance-review.md, smart_photo_node.py, workflows/photo-enhance.json (and empty workflows/ dir), and __pycache__/smart_photo_node.cpython-314.pyc (and empty __pycache__/ dir). No .hyperstack-vm-photo-state.json* state files were present. ComfyUI references in lib/hyperstack/*.rb intentionally left for task T2. --- hyperstack-vm-photo.toml | 68 ---- photo-enhance-review.md | 678 ------------------------------------ photo-enhance.rb | 581 ------------------------------- smart_photo_node.py | 796 ------------------------------------------- workflows/photo-enhance.json | 106 ------ 5 files changed, 2229 deletions(-) delete mode 100644 hyperstack-vm-photo.toml delete mode 100644 photo-enhance-review.md delete mode 100755 photo-enhance.rb delete mode 100644 smart_photo_node.py delete mode 100644 workflows/photo-enhance.json diff --git a/hyperstack-vm-photo.toml b/hyperstack-vm-photo.toml deleted file mode 100644 index d6130b0..0000000 --- a/hyperstack-vm-photo.toml +++ /dev/null @@ -1,68 +0,0 @@ -[auth] -api_key_file = "~/.hyperstack" - -[hyperstack] -base_url = "https://infrahub-api.nexgencloud.com/v1" - -[state] -file = ".hyperstack-vm-photo-state.json" - -[vm] -name_prefix = "hyperstack-photo" -hostname = "hyperstack-photo" -environment_name = "snonux-ollama" - -# L40 (48GB GDDR6, ~$1.00/hr) is the recommended GPU for ComfyUI photo enhancement. -# It provides ample VRAM for SUPIR (FP16 needs ~12-20GB) + Real-ESRGAN with large tile sizes. -# Cheaper than A100 ($2/hr) while offering faster CUDA cores and FP8 support. -flavor_name = "n3-L40x1" -image_name = "Ubuntu Server 24.04 LTS R570 CUDA 12.8 with Docker" -assign_floating_ip = true -create_bootable_volume = false -enable_port_randomization = false -labels = ["comfyui", "photo-enhance", "wireguard"] - -[ssh] -username = "ubuntu" -private_key_path = "~/.ssh/id_rsa" -hyperstack_key_name = "earth" -port = 22 -connect_timeout_sec = 10 - -[network] -wireguard_udp_port = 56710 -wireguard_subnet = "192.168.3.0/24" -# Photo VM uses .4; VM1=.1, VM2=.3 are the LLM VMs on the same wg1 tunnel. -wireguard_server_ip = "192.168.3.4" -# Reuse the same inference port constant; ComfyUI uses its own port 8188 below. -ollama_port = 11434 -allowed_ssh_cidrs = ["auto"] -allowed_wireguard_cidrs = ["auto"] - -[bootstrap] -enable_guest_bootstrap = true -install_wireguard = true -configure_ufw = true -configure_ollama_host = false - -[ollama] -# Not needed on this VM; photo enhancement uses ComfyUI exclusively. -install = false - -[vllm] -# Not needed on this VM; ComfyUI handles all inference. -install = false - -[comfyui] -install = true -# ComfyUI REST API port — opened on the WireGuard subnet only. -port = 8188 -# Model weights on ephemeral NVMe for fast access; survives reboots on Hyperstack. -models_dir = "/ephemeral/comfyui/models" -output_dir = "/ephemeral/comfyui/output" -container_name = "comfyui_photo" -# Pre-downloaded model weights: -# RealESRGAN_x4plus — fast 4x upscaling + sharpening (~65MB, upscale_models/) -# SUPIR-v0Q — SDXL-based photo restoration, Photolemur-quality results (~8GB, checkpoints/) -# SUPIR-v0F — SUPIR variant tuned for faithful fidelity over generative enhancement (~8GB) -models = ["RealESRGAN_x4plus", "SUPIR-v0Q"] diff --git a/photo-enhance-review.md b/photo-enhance-review.md deleted file mode 100644 index 935c0b2..0000000 --- a/photo-enhance-review.md +++ /dev/null @@ -1,678 +0,0 @@ -Main technical risks -1. Real-ESRGAN first, on every image, is your biggest quality risk - -Running every image through: - -4× ESRGAN -then downscale back to original size - -can definitely improve some photos, but it can also introduce: - -hallucinated texture -crispy foliage -waxy skin after interaction with later steps -fake edge detail -zippering around fine geometry -over-defined JPEG blocks on already compressed Fuji JPEGs - -This is the part I would treat as conditionally applied, not universal. - -My recommendation: - -gate ESRGAN based on image characteristics -or at least use different strength paths for portrait vs landscape vs night - -Examples: - -portraits: maybe skip global ESRGAN or use a weaker path -night/high-ISO: be careful, because ESRGAN can turn noise into invented detail -landscapes/architecture: often benefit the most - -Right now the pipeline assumes “seen at 16K then downscaled” is always a win. It often is not. - -2. CodeFormer after global enhancement can amplify inconsistency - -CodeFormer is useful, but it can produce faces that look slightly detached from the rest of the frame if the global pipeline has already altered texture and local contrast. - -Potential issues: - -face crops look cleaner than surrounding skin/neck/hair -restored face sharpness conflicts with depth blur/sharpen later -multiple faces in one frame may get uneven treatment - -Things to consider: - -apply CodeFormer only when face size exceeds a threshold -use a lower-strength/fidelity profile depending on scene -skip CodeFormer for distant faces -log face count and face bounding-box size into metadata - -That would make the workflow easier to debug when faces look “too AI.” - -3. Scene classification using 8 CLIP prompts is clever but brittle - -This is a nice lightweight idea, but it is likely the weakest decision point in the pipeline because eight prompts force coarse categorization. - -Possible failure cases: - -beach sunset might oscillate between beach, golden_hour, and landscape -indoor portraits near a window may flip to portrait or indoor -urban night scenes may misclassify between street and night -cloudy mountain lake might be overcast vs landscape - -Because your grade profile changes exposure/contrast/saturation/detail/denoise, a wrong label can materially alter the image. - -Better approach: - -store the full prompt score distribution, not just argmax -use top-2 or top-3 labels -blend profiles based on confidence instead of hard-switching - -For example: - -60% landscape + 40% golden_hour -instead of forcing one profile - -That would reduce sudden profile mistakes. - -4. CPU image ops at 4K are fine, but not yet optimized as a pipeline - -Your CPU-bound stages are sensible, but there are some efficiency concerns: - -guidedFilter and morphology/blur passes at full 4K are not trivial -ImageScaleBy 16K → 4K on CPU may be heavier than it looks -repeated color-space conversions and full-frame copies can become memory-bandwidth bound -if you later parallelize multiple photos, CPU becomes the bottleneck before GPU memory does - -This matters because your throughput is already 40–50s/photo, and if you batch more aggressively you may saturate host CPU. - -I would especially watch: - -OpenCV allocations -Python ↔ tensor conversion overhead inside custom nodes -whether large intermediate tensors are duplicated unnecessarily -5. Polling /history/ every 2s is workable but not ideal - -It is acceptable, but it is a weak point operationally. - -Risks: - -stale/incomplete history states -long-run prompt ambiguity if ComfyUI restarts -polling delay adds latency -harder recovery when output partially exists but metadata doesn’t - -If ComfyUI or your wrapper supports websocket progress or event-driven status, that would be better. If not, I would at least strengthen state validation: - -ensure expected output files exist and are complete -ensure metadata JSON corresponds to the same prefix -distinguish timeout from partial success -Biggest architectural improvement opportunities -1. Add conditional routing, not one fixed pipeline for every photo - -Right now the graph is elegant, but it is still mostly single-path. - -A more robust system would route based on detected attributes: - -no faces → skip CodeFormer -little/no sky → skip SkyEnhance -low confidence scene label → use default conservative grade -low-detail or noisy photo → reduce or skip ESRGAN -already high-contrast/high-saturation image → apply weaker grade - -That would reduce over-processing and save time. - -2. Move from hardcoded profiles to measured image statistics - -Your scene profiles are sensible, but they are still hand-tuned guesses. - -A stronger next step would be to incorporate measured stats such as: - -luminance histogram -highlight clipping ratio -shadow floor occupancy -saturation percentile -edge density -noise estimate -face area percentage -sky coverage - -Then use those stats to modulate: - -exposure -saturation -detail multiplier -denoise -background blur - -That would make the pipeline more adaptive and less prompt-dependent. - -3. Preserve and restore metadata more deliberately - -You correctly bake orientation before upload because ComfyUI strips EXIF. Good. - -But converting final PNG to JPEG without explicit metadata handling means you may be losing: - -original EXIF fields -capture time -lens/camera info -ICC profile -GPS if present -copyright/author data - -That may be fine, but if the intent is “enhanced derivative of original photo,” I would consider: - -copying selected EXIF fields from source to final JPEG -preserving or explicitly assigning ICC profile -adding software tag / processing note -optionally stripping privacy-sensitive fields by choice, not by accident - -Color profile handling is especially important. “No colour corrections” is not the same as “color managed.” - -4. Add resumability per stage, not just per photo - -Your manifest marks a photo done after full completion, which is good, but partial reruns still require redoing all remote processing for failed photos. - -You could get stronger resilience with stage-aware artifacts: - -oriented temp exists -upload completed -prompt submitted -output downloaded -JPEG written -metadata written - -That might be too much overhead for a personal workflow, but even just logging prompt_id per source photo would help a lot with crash recovery. - -5. Treat JPEG as an output format decision, not a fixed end state - -JPEG quality 92 is reasonable, but for some images: - -foliage -gradients in skies -deep edits after enhancement - -JPEG may reintroduce artifacts after all that expensive work. - -Consider: - -archival output as PNG or TIFF -delivery output as JPEG -optional WebP/AVIF for web usage - -Even if you keep JPEG as primary, having a “master enhanced output” option would be useful. - -Specific comments on the custom stages -AdaptivePhotoGrade - -This is the most promising custom logic in the workflow. - -Good: - -exposure in linear light -contrast and saturation as explicit steps -detail/base decomposition -per-scene profiles - -Concerns: - -gamma 2.2 approximation is simple, but true sRGB transfer is not exactly 2.2 -clipping highlights at 1.0 can lose recoverable rolloff smoothness -HSV saturation edits can behave poorly in skin tones and near highlights -fixed midpoint contrast around 0.5 is simple but not content-aware - -If you keep evolving it, the next quality wins will likely come from: - -proper sRGB transfer functions -luminance-aware saturation -highlight/shadow selective controls -local contrast constrained by noise estimate -SkyEnhance - -Clever and cheap. Good for a CPU stage. - -Risks: - -blue clothing, windows, water, reflective buildings, and tinted glass can get caught -sunset banding or haloing near trees/buildings -vertical prior helps, but can still fail on mountains or upside-weighted compositions - -I would recommend logging: - -sky coverage % -mean mask confidence -whether sky enhancement was effectively skipped - -And maybe auto-disable when coverage is too low or too fragmented. - -DepthSelectiveSharpen - -This is an interesting stage, but also easy to overdo. - -Pros: - -more photographic than simple global sharpening -can add subject separation - -Risks: - -relative depth is not segmentation -hair, glasses, transparent objects, fences, and fine branches can create messy transitions -background blur on an already naturally focused image may look synthetic -blur-plus-sharpen in one stage can produce “smartphone portrait mode” artifacts - -I would strongly consider making this more conservative: - -lower default blur -maybe sharpen foreground only, without explicit background blur -or gate blur by scene type and depth confidence - -For many photos, foreground sharpening alone may be enough. - -Performance review - -Your breakdown is believable. - -The biggest performance cost drivers are probably: - -ESRGAN 4× inference -memory movement around the 16K intermediate -downscale from 16K to 4K -Depth Anything inference - -This means the obvious speed/quality tradeoff lever is: - -reducing or conditionally skipping the 4× path - -That one decision could cut runtime materially. - -If you want better throughput later, likely gains are: - -batch submission queue with bounded concurrency -reuse loaded models across jobs, which ComfyUI already helps with -avoid oversized intermediates when not needed -possibly move some CPU image ops to GPU if they become limiting - -But honestly, for 45 photos, the current runtime is already acceptable. - -Operational review - -This is better than average for reliability, but I would still tighten a few things: - -Add stronger failure modes - -Include distinct handling for: - -upload success but prompt submission failure -prompt accepted but no output file -PNG written but metadata JSON missing -metadata JSON malformed -partial download -ComfyUI restart mid-run -Record more run provenance - -In the metadata or manifest, store: - -prompt_id -workflow hash -custom node git commit or file hash -model file hashes -processing duration per stage if possible - -That will make future tuning much easier. - -Be careful with cache busting - -The random prefix solves ComfyUI cache problems, which is practical. But it also removes some chance of safe reuse. A better long-term version is: - -cache key derived from source hash + workflow hash + relevant params - -Then you can intentionally reuse identical results. - -Best improvements to do next - -If I were prioritizing, I would do these in order: - -make ESRGAN conditional -store CLIP confidence distribution and blend scene profiles -make depth stage more conservative -preserve ICC/selected EXIF metadata -improve failure recovery and provenance logging - -Those would give the best quality/reliability gains without rewriting the whole system. - -Bottom line - -This is a well-designed, thoughtful workflow with a good systems mindset behind it. The main weakness is not architecture — it is over-processing risk from always-on enhancement stages, especially: - -full-image ESRGAN -hard scene-class selection -depth-based blur/sharpen - -So my verdict is: - -architecture: strong -operational design: strong -quality strategy: promising, but should become more conditional and confidence-aware -performance: acceptable for the workload -next maturity step: adaptive routing and stronger provenance - -If you want, I can turn this into a more formal engineering review with sections for “keep / change / risks / recommended roadmap,” or I can propose a v2 pipeline that should improve both quality and speed. - -yes: - -# Photo Enhancement Workflow Review - -## Executive summary - -This workflow is well designed. It has a strong separation between orchestration, inference, deterministic image processing, and human review. The system is practical, understandable, and already close to a small production pipeline for batch photo enhancement. - -The main issue is not the infrastructure or code shape. The main issue is **quality control under an always-on enhancement strategy**. Several expensive stages are applied to every image, even though their benefit is scene-dependent and sometimes negative. The biggest gains now will come from making the workflow **conditional, confidence-aware, and slightly more conservative**. - ---- - -# What should stay - -## 1. Ruby as the control plane - -This is a good choice. - -It gives you: - -* clean batch orchestration -* simple manifest handling -* file lifecycle control -* easy VM lifecycle integration -* a place to keep business logic out of ComfyUI - -## 2. ComfyUI as the execution graph - -Also a good choice. - -It gives you: - -* model reuse -* visual graph structure -* easy injection of runtime parameters -* modular custom node expansion - -## 3. Metadata sidecar generation - -This is one of the strongest parts of the system. - -The `_e.md` and JSON sidecars make the workflow: - -* debuggable -* reviewable -* reproducible -* easier to tune later - -## 4. Human review tool - -The comparison tool is exactly the right final step. Enhancement pipelines often fail because they assume “processed” means “better.” Yours does not. - -## 5. EXIF orientation bake before upload - -Correct and necessary. Good defensive engineering. - ---- - -# What should change - -## 1. Stop treating enhancement as a single fixed path - -Right now the graph is elegant, but too uniform. The workflow should become a **decision tree**, not a single mandatory sequence. - -Some stages should be optional: - -* Real-ESRGAN -* CodeFormer -* SkyEnhance -* DepthSelectiveSharpen -* grading strength inside AdaptivePhotoGrade - -## 2. Make Real-ESRGAN conditional - -This is the highest-priority change. - -Current risks: - -* synthetic texture -* over-crisp foliage -* JPEG artifact amplification -* invented microdetail -* unnatural skin/hair - -### Recommendation - -Use ESRGAN only when: - -* high detail scenes (landscape, architecture) -* strong edge density -* visible softness or compression - -Avoid or weaken for: - -* portraits -* night/high ISO -* already sharp JPEGs - -## 3. Replace hard scene labels with blended grading - -Current approach uses argmax from CLIP. - -Problem: scenes are often mixed. - -### Recommendation - -* keep top 2–3 scene scores -* normalize -* blend profile parameters - -Example: - -* 0.55 landscape -* 0.35 golden_hour -* 0.10 overcast - -Blend exposure, contrast, saturation, detail, denoise. - -## 4. Make depth processing more conservative - -Default behavior should be: - -* foreground sharpening only -* no background blur by default - -Enable blur only when: - -* strong subject separation -* portrait-like composition - -## 5. Preserve metadata intentionally - -Current pipeline likely loses: - -* EXIF -* ICC profile - -### Recommendation - -Preserve or explicitly manage: - -* capture timestamp -* camera/lens info -* ICC profile -* add processing metadata - ---- - -# Main risks - -## Quality risks - -### Over-processing - -Stacked enhancements may lead to synthetic look. - -### Face inconsistency - -CodeFormer may produce mismatch with surrounding regions. - -### Masking errors - -Sky and depth masks may: - -* misclassify regions -* create halos - -## Operational risks - -### Partial success ambiguity - -Need stronger validation for: - -* missing metadata -* partial downloads - -### Weak provenance - -Should log: - -* prompt_id -* workflow hash -* model versions - -### CPU bottleneck - -Potential hotspots: - -* large rescaling -* guided filtering -* morphology operations - ---- - -# Performance review - -## Current state - -~40–50s/photo is acceptable. - -## Main optimization lever - -Make ESRGAN conditional. - -## Secondary lever - -Skip unnecessary stages when not needed. - ---- - -# Recommended v2 architecture - -## Goal - -Make workflow adaptive. - -## Pipeline - -### Stage 0 — Preflight analysis - -Compute: - -* brightness histogram -* saturation -* edge density -* noise estimate -* face stats -* sky coverage -* CLIP scores - -### Stage 1 — Policy selection - -Decide: - -* ESRGAN mode -* CodeFormer usage -* grading blend -* sky enhance on/off -* depth mode - -### Stage 2 — Enhancement - -Run only selected stages. - -### Stage 3 — Output + metadata - -Include: - -* policy decisions -* confidence scores -* timings - ---- - -# Example metadata (v2) - -```json -{ - "workflow_version": "photo-enhance-v2", - "analysis": { - "scene_scores": { - "landscape": 0.51, - "golden_hour": 0.28 - }, - "face_count": 1, - "sky_coverage_pct": 23.4 - }, - "policy": { - "esrgan_mode": "weak", - "depth_mode": "sharpen_only" - } -} -``` - ---- - -# Roadmap - -## Phase 1 - -* conditional ESRGAN -* blended scene grading -* disable background blur default -* preserve metadata - -## Phase 2 - -* preflight analysis -* gating logic for faces and sky -* improved logging - -## Phase 3 - -* better color handling (true sRGB) -* noise-aware detail -* improved saturation logic - ---- - -# Final verdict - -## Strengths - -* strong architecture -* practical workflow -* good separation of concerns - -## Weakness - -* over-processing risk from always-on stages - -## Key improvement - -Move from fixed pipeline → adaptive pipeline - -This will improve both quality and performance significantly. - diff --git a/photo-enhance.rb b/photo-enhance.rb deleted file mode 100755 index 2d6b8de..0000000 --- a/photo-enhance.rb +++ /dev/null @@ -1,581 +0,0 @@ -#!/usr/bin/env ruby -# frozen_string_literal: true - -# photo-enhance.rb — AI photo enhancer via ComfyUI on a Hyperstack GPU VM. -# -# Submits images from --indir to the ComfyUI REST API, downloads the AI-enhanced -# results and saves alongside the originals with an _e suffix. Also downloads -# a per-photo JSON metadata file written by the WritePhotoMetadata ComfyUI node -# and converts it to a human-readable .md report alongside each enhanced photo. -# -# AI pipeline (ComfyUI, GPU): -# 1. Real-ESRGAN realesr-general-x4v3 — 4× upscale at full 4K input, AI denoise -# 2. CodeFormer fidelity=0.7 — neural face restoration -# 3. CLIP ViT-B/32 — scene classification (portrait/landscape/…) -# 4. AdaptivePhotoGrade — scene-tuned exposure/contrast/saturation/detail -# 5. SkyEnhance — HSV sky mask + graduated sky correction -# 6. Depth Anything V2 Small — depth map → foreground sharp, background soft -# -# Usage: -# ruby photo-enhance.rb --config hyperstack-vm-photo.toml \ -# --indir ~/Pictures [--watch] [--workflow workflows/photo-enhance.json] -# -# Requirements: -# - ComfyUI VM: ruby hyperstack.rb --config hyperstack-vm-photo.toml create -# - WireGuard tunnel active (wg1) - -begin - require 'bundler/setup' -rescue LoadError, Gem::GemNotFoundException, Gem::LoadError, Errno::ENOENT - nil -end - -require 'json' -require 'net/http' -require 'optparse' -require 'fileutils' -require 'digest' -require 'time' -require 'set' - -begin - require 'toml-rb' -rescue LoadError - warn "Missing dependency: toml-rb. Run `bundle install` in #{__dir__} first." - exit 2 -end - -# --------------------------------------------------------------------------- -# Config -# --------------------------------------------------------------------------- - -class PhotoConfig - attr_reader :host, :port, :workflow_path - - def initialize(config_path, workflow_path_override) - raw = TomlRB.load_file(File.expand_path(config_path)) - hostname = raw.dig('vm', 'hostname') || 'hyperstack-photo' - interface = raw.dig('local_client', 'interface_name') || 'wg1' - @host = "#{hostname}.#{interface}" - @port = Integer(raw.dig('comfyui', 'port') || 8188) - @workflow_path = workflow_path_override || - File.join(File.dirname(File.expand_path(config_path)), 'workflows', 'photo-enhance.json') - end -end - -# --------------------------------------------------------------------------- -# ComfyUI API client — upload, submit, poll, download. -# --------------------------------------------------------------------------- - -class ComfyUIClient - POLL_INTERVAL_SEC = 2 - POLL_TIMEOUT_SEC = 300 # 5 minutes; ESRGAN is fast on GPU - # When ComfyUI crashes (OOM), systemd restarts it in ~30s. - # We poll until reachable again, up to this many seconds total. - RECOVERY_TIMEOUT_SEC = 300 - RECOVERY_POLL_SEC = 10 - - def initialize(host:, port:, out: $stdout) - @host = host - @port = port - @out = out - end - - def upload_image(file_path) - filename = File.basename(file_path) - image_data = File.binread(file_path) - boundary = "----RubyPhotoEnhance#{hex(8)}" - body = [ - "--#{boundary}\r\n", - "Content-Disposition: form-data; name=\"image\"; filename=\"#{filename}\"\r\n", - "Content-Type: #{mime_type(file_path)}\r\n\r\n", - image_data, - "\r\n--#{boundary}\r\n", - "Content-Disposition: form-data; name=\"overwrite\"\r\n\r\ntrue\r\n", - "--#{boundary}--\r\n" - ].join - resp = post_raw('/upload/image', body, "multipart/form-data; boundary=#{boundary}") - raise "Upload failed (#{resp.code}): #{resp.body}" unless resp.code == '200' - JSON.parse(resp.body)['name'] || filename - rescue Errno::ECONNREFUSED, Errno::EHOSTUNREACH, SocketError => e - raise "Cannot reach ComfyUI at #{@host}:#{@port} — is WireGuard active? (#{e.message})" - end - - def submit_prompt(workflow) - resp = post_json('/prompt', { 'prompt' => workflow }) - raise "Prompt failed (#{resp.code}): #{resp.body}" unless resp.code == '200' - JSON.parse(resp.body)['prompt_id'] or raise "No prompt_id in: #{resp.body}" - end - - def wait_for_output(prompt_id) - deadline = Time.now + POLL_TIMEOUT_SEC - loop do - raise "Timed out after #{POLL_TIMEOUT_SEC}s for #{prompt_id}" if Time.now > deadline - - resp = get("/history/#{prompt_id}") - raise "History poll failed (#{resp.code})" unless resp.code == '200' - - result = JSON.parse(resp.body)[prompt_id] - if result - outputs = extract_filenames(result) - return outputs unless outputs.empty? - - # ComfyUI cached the run (identical inputs) and wrote no new files — bail fast. - status = result.dig('status', 'status_str') - raise "ComfyUI cached execution returned no outputs for #{prompt_id}" \ - if result.dig('status', 'completed') && status == 'success' - end - - sleep POLL_INTERVAL_SEC - end - end - - def download_output(filename, dest_path) - resp = get("/view?filename=#{URI.encode_www_form_component(filename)}&type=output&subfolder=") - raise "Download failed (#{resp.code}) for #{filename}" unless resp.code == '200' - FileUtils.mkdir_p(File.dirname(dest_path)) - File.binwrite(dest_path, resp.body) - end - - def check_connectivity! - resp = get('/system_stats') - raise "Health check failed (#{resp.code}): #{resp.body}" unless resp.code == '200' - rescue Errno::ECONNREFUSED, Errno::EHOSTUNREACH, SocketError => e - raise "Cannot reach ComfyUI at #{@host}:#{@port} — is WireGuard active? (#{e.message})" - end - - # Polls ComfyUI until it responds again (or times out). - # Called automatically when an upload or submit fails with a connection error. - # ComfyUI crashes on OOM (large ESRGAN tensors) and systemd restarts it in ~30s. - # Returns true if recovered, raises on timeout. - def wait_for_recovery - @out.puts " ComfyUI unreachable — waiting for restart (up to #{RECOVERY_TIMEOUT_SEC}s)..." - deadline = Time.now + RECOVERY_TIMEOUT_SEC - start = Time.now - loop do - raise "ComfyUI did not recover within #{RECOVERY_TIMEOUT_SEC}s — giving up" if Time.now > deadline - - sleep RECOVERY_POLL_SEC - begin - get('/system_stats') - @out.puts ' ComfyUI recovered — resuming' - return true - rescue Errno::ECONNREFUSED, Errno::EHOSTUNREACH, SocketError, Net::OpenTimeout - @out.puts " still waiting... (#{(Time.now - start).round}s elapsed)" - end - end - raise "ComfyUI did not recover within #{RECOVERY_TIMEOUT_SEC}s — giving up" - end - - private - - def extract_filenames(result) - Array(result.dig('outputs')) - .flat_map { |_id, node| Array(node['images']) } - .map { |img| img['filename'] } - .compact.reject(&:empty?) - end - - def get(path) - uri = URI("http://#{@host}:#{@port}#{path}") - Net::HTTP.start(uri.host, uri.port, open_timeout: 10, read_timeout: 120) { |h| h.get(uri) } - end - - def post_json(path, payload) - uri = URI("http://#{@host}:#{@port}#{path}") - req = Net::HTTP::Post.new(uri) - req['Content-Type'] = 'application/json' - req.body = JSON.generate(payload) - Net::HTTP.start(uri.host, uri.port, open_timeout: 10, read_timeout: 120) { |h| h.request(req) } - end - - def post_raw(path, body, content_type) - uri = URI("http://#{@host}:#{@port}#{path}") - req = Net::HTTP::Post.new(uri) - req['Content-Type'] = content_type - req.body = body - Net::HTTP.start(uri.host, uri.port, open_timeout: 10, read_timeout: 120) { |h| h.request(req) } - end - - def mime_type(path) - case File.extname(path).downcase - when '.jpg', '.jpeg' then 'image/jpeg' - when '.png' then 'image/png' - when '.webp' then 'image/webp' - else 'application/octet-stream' - end - end - - def hex(n) - Digest::SHA256.hexdigest(Time.now.to_f.to_s + rand.to_s)[0, n * 2] - end -end - -# --------------------------------------------------------------------------- -# Manifest — avoids re-processing files across runs and in watch mode. -# --------------------------------------------------------------------------- - -class ProcessedManifest - FILE_NAME = '.photo-enhance-processed' - - def initialize(dir) - @path = File.join(dir, FILE_NAME) - @entries = load_entries - end - - def processed?(file_path) - @entries.include?(digest(file_path)) - end - - def mark_done(file_path) - key = digest(file_path) - @entries << key - File.open(@path, 'a') { |f| f.puts(key) } - end - - private - - def load_entries - return Set.new unless File.exist?(@path) - File.readlines(@path, chomp: true).map(&:strip).reject(&:empty?).to_set - end - - # Covers basename + size + mtime so a re-shot of the same filename re-processes. - def digest(file_path) - stat = File.stat(file_path) - Digest::SHA256.hexdigest("#{File.basename(file_path)}:#{stat.size}:#{stat.mtime.to_i}") - rescue Errno::ENOENT - Digest::SHA256.hexdigest(File.basename(file_path)) - end -end - -# --------------------------------------------------------------------------- -# Enhancer — orchestrates upload → AI → download → colour correct per image. -# --------------------------------------------------------------------------- - -class PhotoEnhancer - SUPPORTED_EXTENSIONS = %w[.jpg .jpeg .png .webp .raf .cr2 .cr3 .nef .arw .dng .rw2].freeze - RAW_EXTENSIONS = %w[.raf .cr2 .cr3 .nef .arw .dng .rw2].freeze - - # No colour corrections — pure AI output from Real-ESRGAN is used as-is. - # ImageMagick is only used to bake EXIF rotation and convert PNG→JPEG. - COLOR_ARGS = [].freeze - - def initialize(config:, client:, workflow:, indir:, manifest:, out: $stdout) - @config = config - @client = client - @workflow = workflow - @indir = indir - @manifest = manifest - @out = out - end - - def run(watch: false) - @client.check_connectivity! - @out.puts "ComfyUI ready at http://#{@config.host}:#{@config.port}" - @out.puts "Enhancing photos in #{@indir}" - @out.puts watch ? '(watch mode — Ctrl-C to stop)' : '' - - loop do - find_pending.each { |path| enhance_one(path) } - break unless watch - sleep 5 - end - end - - private - - def find_pending - Dir.glob(File.join(@indir, '*')) - .select { |f| File.file?(f) && SUPPORTED_EXTENSIONS.include?(File.extname(f).downcase) } - .reject { |f| File.basename(f, '.*').end_with?('_e') } - .reject { |f| File.basename(f).include?('.orient.') } # temp decode files (.orient.tiff etc.) - .reject { |f| @manifest.processed?(f) } - .sort - end - - def enhance_one(src_path) - ext = File.extname(src_path).downcase - basename = File.basename(src_path, File.extname(src_path)) - # RAW files are always output as JPEG — there is no enhanced RAW format. - out_ext = RAW_EXTENSIONS.include?(ext) ? '.jpg' : ext - dest_path = File.join(File.dirname(src_path), "#{basename}_e#{out_ext}") - - @out.puts "[#{Time.now.strftime('%H:%M:%S')}] #{File.basename(src_path)}" - - # Bake in EXIF rotation before uploading — ComfyUI strips EXIF metadata. - upload_path = auto_orient_tempfile(src_path) - - begin - retried = false - begin - uploaded_name = @client.upload_image(upload_path) - workflow = inject_input(@workflow, uploaded_name) - prompt_id = @client.submit_prompt(workflow) - @out.puts " prompt #{prompt_id}" - - filenames = @client.wait_for_output(prompt_id) - raise "No outputs returned for #{src_path}" if filenames.empty? - rescue RuntimeError => e - # On connection refused (ComfyUI crashed / OOM), wait for systemd to restart - # it and retry this photo once. Any other error propagates immediately. - if !retried && e.message.include?('Cannot reach ComfyUI') - retried = true - @client.wait_for_recovery - retry - end - raise - end - - # ComfyUI outputs PNG; download then convert to output format. - tmp_png = "#{dest_path}.tmp.png" - @client.download_output(filenames.first, tmp_png) - save_with_corrections(tmp_png, dest_path, out_ext) - - # Restore original EXIF metadata onto the enhanced JPEG. - # ComfyUI strips all EXIF when it processes the image; this brings back - # capture time, camera/lens info, ICC profile, and GPS coordinates. - copy_exif(src_path, dest_path) - - # Download the JSON metadata written by WritePhotoMetadata and render it - # as a human-readable .md report alongside the enhanced photo. - # ComfyUI appends _NNNNN_ counter: "enhanced_abc123__00001_.png" → "enhanced_abc123_" - prefix = filenames.first.sub(/_\d+_\.png$/, '') - meta_file = "#{prefix}meta.json" - md_path = File.join(File.dirname(dest_path), - "#{File.basename(dest_path, File.extname(dest_path))}.md") - download_and_write_md(meta_file, src_path, dest_path, md_path, prompt_id) - - @manifest.mark_done(src_path) - @out.puts " -> #{dest_path} (#{kb(src_path)} KB -> #{kb(dest_path)} KB)" - ensure - # Always remove the oriented tempfile (and any downloaded PNG temp) - # so failures do not leave orphaned files on disk. - File.delete(tmp_png) if defined?(tmp_png) && File.exist?(tmp_png) - File.delete(upload_path) if upload_path != src_path && File.exist?(upload_path) - end - rescue StandardError => e - @out.puts " ERROR #{File.basename(src_path)}: #{e.message}" - end - - # Decode and orient the source image into a temp file suitable for ComfyUI upload. - # RAW files (RAF, CR2, NEF…) are decoded to 16-bit TIFF via ImageMagick's LibRaw - # delegate — ComfyUI's LoadImage cannot read RAW formats directly. - # JPEG/PNG inputs are just auto-oriented in place. - # Falls back to the original path if magick is unavailable. - def auto_orient_tempfile(src_path) - ext = File.extname(src_path).downcase - # RAW → TIFF so ComfyUI can load it; JPEG/PNG → same extension with rotation baked in. - out_ext = RAW_EXTENSIONS.include?(ext) ? '.tiff' : ext - tmp = "#{src_path}.orient#{out_ext}" - return tmp if system('magick', src_path, '-auto-orient', tmp) && File.exist?(tmp) - - @out.puts " Warning: auto-orient failed, uploading original" - src_path - end - - # Convert the downloaded PNG to the target format (JPEG quality 92 for .jpg). - # No colour processing — pure AI output from Real-ESRGAN is preserved as-is. - def save_with_corrections(src_png, dest_path, ext) - quality_args = ext.match?(/\.jpe?g/) ? ['-quality', '92'] : [] - system('magick', src_png, *COLOR_ARGS, *quality_args, dest_path) - end - - # Copy selected EXIF fields from the original source file to the enhanced JPEG. - # ComfyUI strips all metadata during inference; this restores capture time, - # camera/lens info, ICC profile, and GPS so the output is a proper derivative. - # Thumbnail and PreviewImage are excluded — they would show the un-enhanced original. - def copy_exif(src_path, dest_path) - return unless dest_path.match?(/\.jpe?g$/i) - return unless system('which', 'exiftool', out: File::NULL, err: File::NULL) - - # Copy all EXIF/IPTC/GPS/ICC tags from source, skipping embedded previews. - # Orientation is excluded because auto-orient already baked rotation into the pixels — - # restoring the original tag would cause viewers to double-rotate the image. - unless system( - 'exiftool', - '-TagsFromFile', src_path, - '-all:all', - '--Orientation', # baked in by magick -auto-orient; don't restore old tag - '--ThumbnailImage', # skip old thumbnail (shows un-enhanced photo) - '--PreviewImage', # skip full preview too - '-overwrite_original', - dest_path, - out: File::NULL, err: File::NULL - ) - @out.puts " Warning: exiftool copy_exif failed for #{File.basename(dest_path)}" - return - end - - # Explicitly set Orientation=1 (normal) so all viewers agree rotation is done. - system('exiftool', '-overwrite_original', '-Orientation=1', '-n', dest_path, - out: File::NULL, err: File::NULL) - - # Tag the output as a derived image so viewers know it was processed - system( - 'exiftool', - '-overwrite_original', - '-Software=photo-enhance (Real-ESRGAN + ComfyUI)', - dest_path, - out: File::NULL, err: File::NULL - ) - end - - # Download the WritePhotoMetadata JSON from ComfyUI output and render it - # as a Markdown report saved alongside the enhanced photo. - # prompt_id is included in the report for reproducibility and crash recovery. - def download_and_write_md(meta_filename, src_path, dest_path, md_path, prompt_id = nil) - resp = @client.send(:get, - "/view?filename=#{URI.encode_www_form_component(meta_filename)}&type=output&subfolder=") - return unless resp.code == '200' - - meta = JSON.parse(resp.body) - profile = meta['enhancement_profile'] || {} - sky = meta['sky'] || {} - depth = meta['depth_sharpen'] || {} - models = meta['models'] || {} - scene = meta['scene_type'] || 'unknown' - esrgan_mode = meta['esrgan_mode'] || 'full' - scene_scores = meta['scene_scores'] || {} - ts = meta['generated_at'] || Time.now.utc.iso8601 - - # Format top-3 scene confidence scores as "landscape 55%, golden_hour 35%, overcast 10%" - scores_str = scene_scores - .sort_by { |_, v| -v } - .map { |s, v| "#{s} #{(v * 100).round}%" } - .join(', ') - - md = <<~MD - # #{File.basename(dest_path)} — Enhancement Report - - **Source:** #{File.basename(src_path)} (#{kb(src_path)} KB) - **Enhanced:** #{File.basename(dest_path)} (#{kb(dest_path)} KB) - **Processed:** #{ts} - **ComfyUI prompt ID:** #{prompt_id || 'n/a'} - - ## AI Pipeline - - | Step | Model / Node | Device | What it does | - |------|-------------|--------|--------------| - | 1 | `#{models['scene_detect']}` | GPU | Zero-shot scene classification → ESRGAN gating | - | 2 | `#{models['upscaler']}` (#{esrgan_mode}) | GPU | 4× upscale at full 4K input → 16K → back to 4K | - | 3 | `#{models['face_restore']}` | GPU | Face detection + neural restoration | - | 4 | Adaptive Photo Grade | CPU | Scene-blended exposure / contrast / saturation / detail | - | 5 | Sky Enhance | CPU | HSV sky mask + graduated sky correction | - | 6 | `#{models['depth']}` | GPU | Depth map → foreground sharpening | - - ## Scene Detection - - | | | - |-|-| - | **Detected scene** | #{scene} | - | **Top-3 scores** | #{scores_str.empty? ? 'n/a' : scores_str} | - | **ESRGAN mode** | #{esrgan_mode} (skip=portrait/night, weak=indoor/golden, full=landscape/beach) | - - ## Colour Grading Profile (blended) - - | Setting | Value | - |---------|-------| - | Exposure | +#{profile['exposure_stops']} stops | - | Contrast | #{profile['contrast_factor']}× | - | Saturation | #{profile['saturation_mult']}× | - | Detail / Clarity | #{profile['detail_mult']}× | - | Denoise strength | #{profile['denoise_strength']} | - - ## Sky Enhancement - - | Setting | Value | - |---------|-------| - | Sky coverage | #{sky['coverage_pct']}% of image | - | Sky exposure | +#{sky['sky_exposure']} stops | - | Sky saturation | #{sky['sky_saturation']}× | - - ## Depth-Guided Sharpening - - | Setting | Value | - |---------|-------| - | Foreground sharpening | #{depth['foreground_sharpen']}× | - | Background blur | #{depth['background_blur']} (0.0 = disabled) | - MD - - File.write(md_path, md) - rescue StandardError => e - @out.puts " Warning: could not write metadata report: #{e.message}" - end - - # Inject the upload filename and a unique prefix into LoadImage, SaveImage, - # and WritePhotoMetadata to bust ComfyUI's cache and link metadata to image. - def inject_input(workflow, filename) - wf = JSON.parse(JSON.generate(workflow)) # deep dup - prefix = "enhanced_#{Digest::SHA256.hexdigest(Time.now.to_f.to_s + rand.to_s)[0, 8]}_" - wf.each_value do |node| - next unless node.is_a?(Hash) - case node['class_type'] - when 'LoadImage' then node['inputs']['image'] = filename - when 'SaveImage' then node['inputs']['filename_prefix'] = prefix - when 'WritePhotoMetadata' - node['inputs']['filename_prefix'] = prefix - node['inputs']['source_filename'] = filename - end - end - wf - end - - def kb(path) - (File.size(path) / 1024.0).round - end -end - -# --------------------------------------------------------------------------- -# CLI -# --------------------------------------------------------------------------- - -options = { - config: File.join(__dir__, 'hyperstack-vm-photo.toml'), - indir: nil, - watch: false, - workflow: nil, - test: false -} - -OptionParser.new do |o| - o.banner = 'Usage: ruby photo-enhance.rb [options]' - o.on('--config PATH', 'TOML config (default: hyperstack-vm-photo.toml)') { |v| options[:config] = v } - o.on('--indir PATH', 'Directory of photos to enhance') { |v| options[:indir] = v } - o.on('--workflow PATH', 'ComfyUI workflow JSON override') { |v| options[:workflow] = v } - o.on('--watch', 'Keep running, process new images as they arrive') { options[:watch] = true } - o.on('--test', 'Check ComfyUI connectivity only, then exit') { options[:test] = true } - o.on('-h', '--help', 'Show this help') { puts o; exit } -end.parse! - -abort "Config not found: #{options[:config]}" unless File.exist?(options[:config]) - -cfg = PhotoConfig.new(options[:config], options[:workflow]) -client = ComfyUIClient.new(host: cfg.host, port: cfg.port) - -if options[:test] - begin - client.check_connectivity! - puts "ComfyUI reachable at http://#{cfg.host}:#{cfg.port} — OK" - exit 0 - rescue RuntimeError => e - warn "ERROR: #{e.message}"; exit 1 - end -end - -abort '--indir is required' unless options[:indir] -indir = File.expand_path(options[:indir]) -abort "Directory not found: #{indir}" unless File.directory?(indir) -abort "Workflow not found: #{cfg.workflow_path}" unless File.exist?(cfg.workflow_path) - -workflow = JSON.parse(File.read(cfg.workflow_path)) -manifest = ProcessedManifest.new(indir) -enhancer = PhotoEnhancer.new(config: cfg, client: client, workflow: workflow, - indir: indir, manifest: manifest) -begin - enhancer.run(watch: options[:watch]) -rescue RuntimeError => e - warn "ERROR: #{e.message}"; exit 1 -rescue Interrupt - puts "\nStopped." -end diff --git a/smart_photo_node.py b/smart_photo_node.py deleted file mode 100644 index aeb1a3b..0000000 --- a/smart_photo_node.py +++ /dev/null @@ -1,796 +0,0 @@ -""" -Smart Photo Enhancement Nodes for ComfyUI -========================================== -Six AI-driven nodes that replace static colour-correction filters with -content-aware, adaptive processing: - - CLIPSceneDetect — CLIP zero-shot classification → scene label + - top-3 confidence scores as a JSON string - ConditionalESRGANBlend — scene-aware ESRGAN gating: blends original with - the upscaled result at a scene-tuned ratio - (portrait/night skip ESRGAN; landscape/beach keep full) - AdaptivePhotoGrade — scene-tuned exposure/contrast/saturation/detail; - blends multiple scene profiles weighted by CLIP scores - SkyEnhance — HSV sky mask + graduated exposure & saturation boost - DepthSelectiveSharpen — Depth-Anything depth map → foreground sharpening only - (no background blur by default — avoids portrait-mode look) - WritePhotoMetadata — per-photo JSON report: scene scores, ESRGAN mode, - grading profile, sky coverage, depth settings, models - -All heavy models are loaded once and kept in _MODEL_CACHE between prompts. -""" - -import json -import re -import torch -import numpy as np -import cv2 -from PIL import Image - -# --------------------------------------------------------------------------- -# Global model cache — prevents reloading 100–600 MB models every frame -# --------------------------------------------------------------------------- -_MODEL_CACHE: dict = {} - - -def _cached_model(key: str, loader_fn): - """Return a cached model, loading it on the first call.""" - if key not in _MODEL_CACHE: - _MODEL_CACHE[key] = loader_fn() - return _MODEL_CACHE[key] - - -# --------------------------------------------------------------------------- -# CLIPSceneDetect -# --------------------------------------------------------------------------- -class CLIPSceneDetect: - """ - Zero-shot scene classification using OpenAI CLIP (ViT-B/32, ~600 MB). - - Matches the photo against 8 descriptive text prompts via cosine similarity - and emits: - • scene_type — winning scene label (STRING) for downstream nodes - • scene_scores — top-3 scene scores as a normalised JSON string - e.g. '{"landscape": 0.55, "golden_hour": 0.35, "overcast": 0.10}' - - Runs on the ORIGINAL image before ESRGAN so the score can gate upscaling - in ConditionalESRGANBlend. - - Scene labels: portrait | landscape | night | indoor | - golden_hour | overcast | beach | street - """ - - # Text prompts whose cosine similarity to the image selects the scene - SCENE_PROMPTS = [ - "a portrait photograph of a person or people", - "a landscape photograph of nature or scenery outdoors", - "a night photograph taken in low light or darkness", - "an indoor photograph inside a room or building", - "a golden hour or sunset photograph with warm orange light", - "an overcast or cloudy day outdoor photograph", - "a beach, ocean, or waterfront photograph", - "a street, city, or urban photograph", - ] - SCENE_LABELS = [ - "portrait", "landscape", "night", "indoor", - "golden_hour", "overcast", "beach", "street", - ] - - @classmethod - def INPUT_TYPES(cls): - return {"required": {"image": ("IMAGE",)}} - - RETURN_TYPES = ("IMAGE", "STRING", "STRING") - RETURN_NAMES = ("image", "scene_type", "scene_scores") - FUNCTION = "detect" - CATEGORY = "image/smart" - - def detect(self, image): - from transformers import CLIPProcessor, CLIPModel - - device = "cuda" if torch.cuda.is_available() else "cpu" - - def _load(): - print("[CLIPSceneDetect] Loading CLIP ViT-B/32…") - m = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device).eval() - p = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") - return m, p - - model, processor = _cached_model("clip_scene", _load) - - # Use the first image in the batch; all frames are the same scene - img_np = (image[0].cpu().numpy() * 255).astype(np.uint8) - img_pil = Image.fromarray(img_np) - - inputs = processor( - text=self.SCENE_PROMPTS, - images=img_pil, - return_tensors="pt", - padding=True, - ).to(device) - - with torch.no_grad(): - logits = model(**inputs).logits_per_image[0] - probs = logits.softmax(dim=0).cpu() - - # Winning scene for hard-switch compatibility - idx = int(probs.argmax()) - scene = self.SCENE_LABELS[idx] - conf = float(probs[idx]) - print(f"[CLIPSceneDetect] → {scene} ({conf:.1%})") - - # Build top-3 normalised scores for blending downstream nodes. - # Normalising to sum=1.0 lets downstream code weight-blend without - # worrying about softmax temperature or total mass. - top3_idx = probs.topk(3).indices.tolist() - top3_sum = sum(float(probs[i]) for i in top3_idx) - top3_scores = { - self.SCENE_LABELS[i]: round(float(probs[i]) / top3_sum, 4) - for i in top3_idx - } - scene_scores = json.dumps(top3_scores) - print(f"[CLIPSceneDetect] Top-3 scores: {top3_scores}") - - return (image, scene, scene_scores) - - -# --------------------------------------------------------------------------- -# ConditionalESRGANBlend -# --------------------------------------------------------------------------- -class ConditionalESRGANBlend: - """ - Scene-aware ESRGAN gating via pixel-level blending. - - Real-ESRGAN 4× upscale is expensive (~15–20 s/image) and introduces: - • synthetic texture and crispy foliage on already-sharp images - • waxy skin and invented micro-detail on portraits - • amplified JPEG block artifacts on high-ISO night shots - - This node blends the ESRGAN output with the original at a scene-tuned - ratio so over-processed images revert to a safer look without removing - the upscale step from the ComfyUI graph entirely. - - Blend ratios (0.0 = full original, 1.0 = full ESRGAN): - portrait / night → 0.0 (skip — never benefits) - indoor → 0.25 (light touch for interior textures) - golden_hour → 0.60 (moderate — warm skin tones are sensitive) - overcast / street → 0.75 (strong but not maximum) - landscape / beach → 1.0 (full — these gain the most from ESRGAN) - default → 1.0 - - When scene_scores (JSON) are provided, the effective ratio is a - confidence-weighted average of the per-scene ratios — e.g. a 55/35/10 - split between landscape, golden_hour, and overcast yields ~0.86 instead - of the hard 1.0 for landscape alone. - - The esrgan_mode output string ("skip" / "weak" / "full") is written into - the metadata sidecar for debugging. - """ - - # Per-scene blend ratios: 0.0 = original, 1.0 = pure ESRGAN - ESRGAN_RATIOS = { - "portrait": 0.0, - "night": 0.0, - "indoor": 0.25, - "golden_hour": 0.60, - "overcast": 0.75, - "street": 0.75, - "landscape": 1.0, - "beach": 1.0, - "default": 1.0, - } - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "original": ("IMAGE",), - "esrgan": ("IMAGE",), - "scene_type": ("STRING", {"default": "default"}), - "scene_scores": ("STRING", {"default": "{}"}), - } - } - - RETURN_TYPES = ("IMAGE", "STRING") - RETURN_NAMES = ("image", "esrgan_mode") - FUNCTION = "blend" - CATEGORY = "image/smart" - - def blend(self, original, esrgan, scene_type: str, scene_scores: str): - ratio = self._compute_ratio(scene_type, scene_scores) - - if ratio <= 0.02: - esrgan_mode = "skip" - elif ratio >= 0.98: - esrgan_mode = "full" - else: - esrgan_mode = "weak" - - print(f"[ConditionalESRGANBlend] scene={scene_type} ratio={ratio:.2f} mode={esrgan_mode}") - - if esrgan_mode == "skip": - return (original, esrgan_mode) - if esrgan_mode == "full": - return (esrgan, esrgan_mode) - - # Pixel-level blend: result = original*(1-r) + esrgan*r - # Both tensors are [B, H, W, C] float32 0..1 - blended = original * (1.0 - ratio) + esrgan * ratio - return (torch.clamp(blended, 0, 1), esrgan_mode) - - def _compute_ratio(self, scene_type: str, scene_scores_json: str) -> float: - """ - Return the ESRGAN blend ratio. When scene_scores carries top-3 - confidence values, returns a weighted average of per-scene ratios. - Falls back to the hard per-scene ratio if parsing fails. - """ - base_ratio = self.ESRGAN_RATIOS.get(scene_type, self.ESRGAN_RATIOS["default"]) - - try: - scores = json.loads(scene_scores_json) - if not scores: - return base_ratio - except Exception: - return base_ratio - - # Confidence-weighted blend of ratios across the top-3 scenes - total_weight = sum(scores.values()) - if total_weight < 1e-6: - return base_ratio - - weighted_ratio = sum( - self.ESRGAN_RATIOS.get(s, self.ESRGAN_RATIOS["default"]) * w - for s, w in scores.items() - ) / total_weight - return float(weighted_ratio) - - -# --------------------------------------------------------------------------- -# AdaptivePhotoGrade -# --------------------------------------------------------------------------- -class AdaptivePhotoGrade: - """ - Scene-adaptive colour grading node. - - Applies exposure correction (linear-light), contrast, saturation, and - guided-filter clarity enhancement with parameters tuned per scene type. - - When scene_scores (JSON) from CLIPSceneDetect are available, the grading - parameters are blended across the top-3 scene profiles weighted by their - confidence values — e.g. a 55/35/10 landscape/golden_hour/overcast split - produces a weighted average of those three profiles, avoiding the hard - cut artefacts caused by a single misclassified label. - - Falls back to the balanced 'default' profile for unknown scene labels or - when scene_scores is empty/unparseable. - """ - - # Per-scene profiles: exposure in stops, contrast factor, saturation - # multiplier, detail enhancement multiplier, denoise strength (0..1). - PROFILES = { - # Portraits: gentle — preserve skin tones, avoid over-sharpening hair - "portrait": dict(stops=0.30, contrast=1.00, saturation=1.00, detail=1.2, denoise=0.15), - # Landscapes: vivid — strong clarity, saturated skies & greens - "landscape": dict(stops=0.20, contrast=1.05, saturation=1.15, detail=1.8, denoise=0.05), - # Night: lift shadows aggressively, reduce sharpening (hides noise) - "night": dict(stops=0.80, contrast=1.00, saturation=0.90, detail=0.8, denoise=0.30), - # Indoor: correct typically warm/dim ambient light - "indoor": dict(stops=0.50, contrast=1.00, saturation=1.05, detail=1.3, denoise=0.10), - # Golden hour: enhance warmth, lift shadow detail - "golden_hour": dict(stops=0.25, contrast=1.05, saturation=1.20, detail=1.5, denoise=0.05), - # Overcast: punch contrast to compensate for flat light - "overcast": dict(stops=0.40, contrast=1.05, saturation=1.10, detail=1.6, denoise=0.08), - # Beach: bright scene, protect highlights, boost blues/greens - "beach": dict(stops=0.15, contrast=1.00, saturation=1.20, detail=1.7, denoise=0.05), - # Street: punchy contrast, neutral colour - "street": dict(stops=0.35, contrast=1.05, saturation=1.05, detail=1.5, denoise=0.08), - # Balanced fallback for unrecognised labels - "default": dict(stops=0.40, contrast=1.00, saturation=1.05, detail=1.5, denoise=0.10), - } - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "images": ("IMAGE",), - "scene_type": ("STRING", {"default": "default"}), - "scene_scores": ("STRING", {"default": "{}"}), - } - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "grade" - CATEGORY = "image/smart" - - def grade(self, images, scene_type: str, scene_scores: str = "{}"): - p = self._resolve_profile(scene_type, scene_scores) - print(f"[AdaptivePhotoGrade] Scene={scene_type} → {p}") - - results = [] - for img in images: - arr = img.cpu().numpy().copy() # [H, W, C] float32 0..1 - arr = self._apply_exposure(arr, p["stops"]) - arr = self._apply_contrast(arr, p["contrast"]) - arr = self._apply_saturation(arr, p["saturation"]) - arr = self._apply_detail(arr, p["detail"], p["denoise"]) - results.append(torch.from_numpy(arr.clip(0, 1)).float()) - - return (torch.stack(results),) - - def _resolve_profile(self, scene_type: str, scene_scores_json: str) -> dict: - """ - Return a grading profile by blending top-3 scene profiles weighted by - their CLIP confidence scores. Falls back to hard scene_type lookup when - the JSON is absent or malformed. - """ - try: - scores = json.loads(scene_scores_json) - if not scores: - raise ValueError("empty scores") - except Exception: - return self.PROFILES.get(scene_type, self.PROFILES["default"]) - - param_keys = list(next(iter(self.PROFILES.values())).keys()) - total_weight = sum(scores.values()) - if total_weight < 1e-6: - return self.PROFILES.get(scene_type, self.PROFILES["default"]) - - blended = {k: 0.0 for k in param_keys} - for scene, weight in scores.items(): - profile = self.PROFILES.get(scene, self.PROFILES["default"]) - for k in param_keys: - blended[k] += profile[k] * (weight / total_weight) - return blended - - # -- helpers ------------------------------------------------------------ - - def _apply_exposure(self, img: np.ndarray, stops: float) -> np.ndarray: - """ - Per-stop exposure adjustment in linear light. - Converts sRGB → linear, multiplies by 2^stops, clips highlights, converts back. - Simple and photographic — avoids Reinhard's tonal compression which - would darken already-bright Fuji photos. - """ - linear = img ** 2.2 # sRGB → approximate linear - linear = linear * (2.0 ** stops) # shift by N stops (positive = brighter) - return np.clip(linear ** (1.0 / 2.2), 0, 1) # back to sRGB, clip overexposed - - def _apply_contrast(self, img: np.ndarray, factor: float) -> np.ndarray: - """Simple linear contrast around 0.5 midpoint.""" - return np.clip((img - 0.5) * factor + 0.5, 0, 1) - - def _apply_saturation(self, img: np.ndarray, factor: float) -> np.ndarray: - """HSV saturation boost; factor=1.0 is a no-op.""" - u8 = (img * 255).astype(np.uint8) - hsv = cv2.cvtColor(u8, cv2.COLOR_RGB2HSV).astype(np.float32) - hsv[:, :, 1] = np.clip(hsv[:, :, 1] * factor, 0, 255) - return cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2RGB).astype(np.float32) / 255.0 - - def _apply_detail(self, img: np.ndarray, mult: float, denoise: float) -> np.ndarray: - """ - Clarity / structure boost via edge-preserving decomposition. - - Uses cv2.ximgproc.guidedFilter when opencv-contrib is available - (provides the best edge-preserving base layer separation). - Falls back to a bilateral filter base when ximgproc is absent. - - Separates base (low-freq) from detail (high-freq), scales detail by - mult, optionally denoises the base layer via bilateral filter. - """ - u8 = (img * 255).astype(np.uint8) - - # Prefer guided filter (opencv-contrib); fall back to bilateral - try: - base = cv2.ximgproc.guidedFilter(u8, u8, radius=8, eps=int(0.01 * 255 ** 2)) - except AttributeError: - # opencv-contrib not installed — bilateral filter gives a similar - # edge-preserving smooth base at slightly lower quality - sigma = max(15, int(denoise * 75)) - base = cv2.bilateralFilter(u8, d=9, sigmaColor=sigma, sigmaSpace=sigma) - - detail = u8.astype(np.float32) - base.astype(np.float32) - - # Optionally soften the base to reduce noise before adding detail back - if denoise > 0.05: - sigma = int(denoise * 75) - base = cv2.bilateralFilter(base, d=5, sigmaColor=sigma, sigmaSpace=sigma) - - enhanced = base.astype(np.float32) + detail * mult - return np.clip(enhanced / 255.0, 0, 1) - - -# --------------------------------------------------------------------------- -# SkyEnhance -# --------------------------------------------------------------------------- -class SkyEnhance: - """ - Sky region detection and graduated enhancement — no ML model required. - - Detects sky using HSV colour ranges (blue sky, white clouds, sunset tones) - combined with a spatial prior (sky lives in the upper portion of the frame). - Applies independent exposure + saturation adjustments to the sky mask, - blended smoothly with the rest of the image. - - Works on any outdoor shot; portraits and indoor shots receive no change - because the sky mask will be near zero. - """ - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "images": ("IMAGE",), - "sky_exposure": ("FLOAT", {"default": 0.30, "min": -1.0, "max": 1.0, "step": 0.05}), - "sky_saturation": ("FLOAT", {"default": 1.20, "min": 0.5, "max": 2.0, "step": 0.05}), - } - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "enhance" - CATEGORY = "image/smart" - - def enhance(self, images, sky_exposure: float = 0.30, sky_saturation: float = 1.20): - results = [] - for img in images: - arr = (img.cpu().numpy() * 255).astype(np.uint8) - mask = self._detect_sky(arr) - enhanced = self._apply_sky(arr, mask, sky_exposure, sky_saturation) - results.append(torch.from_numpy(enhanced.astype(np.float32) / 255.0)) - return (torch.stack(results),) - - def _detect_sky(self, img_rgb: np.ndarray) -> np.ndarray: - """ - Build a soft float sky mask [0..1] using three HSV colour bands - plus a vertical spatial prior (sky = upper image region). - """ - h = img_rgb.shape[0] - hsv = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2HSV).astype(np.float32) - H, S, V = hsv[:, :, 0], hsv[:, :, 1], hsv[:, :, 2] - - # Band 1: Blue daytime sky (hue 90–140 in OpenCV 0–180 scale) - blue = ((H >= 90) & (H <= 140) & (S >= 30) & (V >= 50)).astype(np.float32) - - # Band 2: White/grey clouds (low saturation, bright) - clouds = ((S < 40) & (V >= 180)).astype(np.float32) - - # Band 3: Sunset/golden sky (hue 0–25 or 155–180, moderate sat) - sunset = (((H <= 25) | (H >= 155)) & (S >= 40) & (V >= 100)).astype(np.float32) - - raw = np.clip(blue + clouds + sunset, 0, 1) - - # Vertical gradient prior: top row = 1.2, bottom row = 0.0 - y_weight = np.linspace(1.2, 0.0, h)[:, np.newaxis] - raw = raw * y_weight - - # Morphological close to fill gaps between cloud patches - kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15)) - raw = cv2.morphologyEx(raw, cv2.MORPH_CLOSE, kernel) - - # Large Gaussian blur for smooth mask edges (avoids halo artifacts) - mask = cv2.GaussianBlur(raw, (51, 51), 0) - return np.clip(mask, 0, 1) - - def _apply_sky(self, img