mirror of
https://github.com/vladmandic/automatic
synced 2026-09-20 01:31:13 +02:00
Merge pull request #4694 from vladmandic/feat/civitai-browser-overhaul
Feat/civitai browser overhaul
This commit is contained in:
@@ -113,6 +113,11 @@ def search_civitai(
|
||||
return []
|
||||
|
||||
t0 = time.time()
|
||||
import re
|
||||
url_match = re.match(r'https?://civitai\.com/models/(\d+)', query.strip())
|
||||
if url_match:
|
||||
query = url_match.group(1)
|
||||
log.info(f'CivitAI: extracted model id={query} from URL')
|
||||
dct = { 'query': query }
|
||||
if len(tag) > 0:
|
||||
dct['tag'] = tag
|
||||
|
||||
+38
-8
@@ -11,6 +11,9 @@ String.prototype.format = function (args) { // eslint-disable-line no-extend-nat
|
||||
let selectedURL = '';
|
||||
let selectedName = '';
|
||||
let selectedType = '';
|
||||
let selectedBase = '';
|
||||
let selectedModelId = '';
|
||||
let selectedVersionId = '';
|
||||
|
||||
function clearModelDetails() {
|
||||
const el = gradioApp().getElementById('model-details') || gradioApp().getElementById('civitai_models_output') || gradioApp().getElementById('models_outcome');
|
||||
@@ -84,7 +87,7 @@ async function modelCardClick(id) {
|
||||
data = data[0]; // assuming the first item is the one we want
|
||||
|
||||
const versionsHTML = data.versions.map((v) => modelVersionsHTML.format({
|
||||
url: `<div class="link" onclick="startCivitDownload('${v.files[0]?.url}', '${v.files[0]?.name}', '${data.type}')"> \udb80\uddda </div>`,
|
||||
url: `<div class="link" onclick="startCivitDownload('${v.files[0]?.url}', '${v.files[0]?.name}', '${data.type}', '${v.base || ''}', ${data.id}, ${v.id})"> \udb80\uddda </div>`,
|
||||
name: v.name || 'unknown',
|
||||
type: v.files[0]?.type || 'unknown',
|
||||
base: v.base || 'unknown',
|
||||
@@ -113,11 +116,14 @@ async function modelCardClick(id) {
|
||||
el.innerHTML = modelHTML;
|
||||
}
|
||||
|
||||
function startCivitDownload(url, name, type) {
|
||||
log('startCivitDownload', { url, name, type });
|
||||
function startCivitDownload(url, name, type, base, modelId, versionId) {
|
||||
log('startCivitDownload', { url, name, type, base, modelId, versionId });
|
||||
selectedURL = [url];
|
||||
selectedName = [name];
|
||||
selectedType = [type];
|
||||
selectedBase = [base || ''];
|
||||
selectedModelId = [modelId || 0];
|
||||
selectedVersionId = [versionId || 0];
|
||||
const civitDownloadBtn = gradioApp().getElementById('civitai_download_btn');
|
||||
if (civitDownloadBtn) civitDownloadBtn.click();
|
||||
}
|
||||
@@ -128,20 +134,44 @@ function startCivitAllDownload(evt) {
|
||||
selectedURL = [];
|
||||
selectedName = [];
|
||||
selectedType = [];
|
||||
selectedBase = [];
|
||||
selectedModelId = [];
|
||||
selectedVersionId = [];
|
||||
for (const version of versions) {
|
||||
const parsed = version.querySelector('td:nth-child(1) div')?.getAttribute('onclick')?.match(/startCivitDownload\('([^']+)', '([^']+)', '([^']+)'\)/);
|
||||
if (!parsed || parsed.length < 4) continue;
|
||||
const parsed = version.querySelector('td:nth-child(1) div')?.getAttribute('onclick')?.match(/startCivitDownload\('([^']+)', '([^']+)', '([^']+)', '([^']*)', (\d+), (\d+)\)/);
|
||||
if (!parsed || parsed.length < 7) continue;
|
||||
selectedURL.push(parsed[1]);
|
||||
selectedName.push(parsed[2]);
|
||||
selectedType.push(parsed[3]);
|
||||
selectedBase.push(parsed[4]);
|
||||
selectedModelId.push(parseInt(parsed[5], 10));
|
||||
selectedVersionId.push(parseInt(parsed[6], 10));
|
||||
}
|
||||
const civitDownloadBtn = gradioApp().getElementById('civitai_download_btn');
|
||||
if (civitDownloadBtn) civitDownloadBtn.click();
|
||||
}
|
||||
|
||||
function downloadCivitModel(modelUrl, modelName, modelType, modelPath, civitToken, innerHTML) {
|
||||
log('downloadCivitModel', { modelUrl, modelName, modelType, modelPath, civitToken });
|
||||
function downloadCivitModel(modelUrl, modelName, modelType, modelBase, mId, vId, modelPath, civitToken, innerHTML) {
|
||||
log('downloadCivitModel', { modelUrl, modelName, modelType, modelBase, mId, vId, modelPath, civitToken });
|
||||
const el = gradioApp().getElementById('civitai_models_output') || gradioApp().getElementById('models_outcome');
|
||||
const currentHTML = el?.innerHTML || '';
|
||||
return [selectedURL, selectedName, selectedType, modelPath, civitToken, currentHTML];
|
||||
return [selectedURL, selectedName, selectedType, selectedBase, selectedModelId, selectedVersionId, modelPath, civitToken, currentHTML];
|
||||
}
|
||||
|
||||
let civitMutualExcludeBound = false;
|
||||
|
||||
function civitaiMutualExclude() {
|
||||
if (civitMutualExcludeBound) return;
|
||||
const searchEl = gradioApp().querySelector('#civit_search_text textarea');
|
||||
const tagEl = gradioApp().querySelector('#civit_search_tag textarea');
|
||||
if (!searchEl || !tagEl) return;
|
||||
civitMutualExcludeBound = true;
|
||||
searchEl.addEventListener('input', () => {
|
||||
tagEl.closest('.gradio-textbox')?.classList.toggle('disabled-look', !!searchEl.value.trim());
|
||||
});
|
||||
tagEl.addEventListener('input', () => {
|
||||
searchEl.closest('.gradio-textbox')?.classList.toggle('disabled-look', !!tagEl.value.trim());
|
||||
});
|
||||
}
|
||||
|
||||
onUiLoaded(civitaiMutualExclude);
|
||||
|
||||
@@ -2236,6 +2236,11 @@ div:has(>#tab-gallery-folders) {
|
||||
filter: blur(0);
|
||||
}
|
||||
|
||||
.disabled-look textarea {
|
||||
opacity: 0.4;
|
||||
pointer-events: none;
|
||||
}
|
||||
|
||||
@keyframes spin {
|
||||
from {
|
||||
transform: rotate(0deg);
|
||||
|
||||
@@ -406,14 +406,24 @@ def download_civit_preview(model_path: str, preview_url: str):
|
||||
return 200, str(total_size), ''
|
||||
|
||||
|
||||
def download_civit_model(model_url: str, model_name: str = '', model_path: str = '', model_type: str = '', token: str = None):
|
||||
def download_civit_model(model_url: str, model_name: str = '', model_path: str = '', model_type: str = '', token: str = None,
|
||||
base_model: str = '', model_id: int = 0, version_id: int = 0):
|
||||
"""Legacy function — delegates to DownloadManager for non-blocking downloads."""
|
||||
if not model_url:
|
||||
log.error('Model download: no url provided')
|
||||
return None
|
||||
if not version_id:
|
||||
import re
|
||||
match = re.search(r'/api/download/models/(\d+)', model_url)
|
||||
if match:
|
||||
version_id = int(match.group(1))
|
||||
from modules.civitai.filemanage_civitai import get_type_folder
|
||||
if not model_path:
|
||||
folder = str(get_type_folder(model_type or 'Checkpoint'))
|
||||
if getattr(shared.opts, 'civitai_save_subfolder_enabled', False):
|
||||
from modules.civitai.filemanage_civitai import resolve_save_path
|
||||
folder = str(resolve_save_path(model_type or 'Checkpoint', model_name=model_name, base_model=base_model))
|
||||
else:
|
||||
folder = str(get_type_folder(model_type or 'Checkpoint'))
|
||||
elif os.path.isabs(model_path):
|
||||
folder = model_path
|
||||
else:
|
||||
@@ -424,6 +434,8 @@ def download_civit_model(model_url: str, model_name: str = '', model_path: str =
|
||||
filename=model_name or "Unknown",
|
||||
model_type=model_type,
|
||||
token=token,
|
||||
model_id=model_id,
|
||||
version_id=version_id,
|
||||
)
|
||||
# Wait for completion (legacy blocking behavior)
|
||||
while item.status in ("queued", "downloading", "verifying"):
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import re
|
||||
import time
|
||||
from installer import log
|
||||
from modules.civitai.client_civitai import client
|
||||
@@ -20,12 +21,18 @@ def search_civitai(
|
||||
token: str = None,
|
||||
exact: bool = True,
|
||||
) -> list[CivitModel]:
|
||||
if not query:
|
||||
log.error('CivitAI: empty query')
|
||||
if not query and not tag and not sort:
|
||||
log.error('CivitAI: no search criteria provided')
|
||||
return []
|
||||
|
||||
t0 = time.time()
|
||||
|
||||
# URL query → extract model ID (e.g. https://civitai.com/models/967405/nova-orange-xl)
|
||||
url_match = re.match(r'https?://civitai\.com/models/(\d+)', query.strip())
|
||||
if url_match:
|
||||
query = url_match.group(1)
|
||||
log.info(f'CivitAI: extracted model id={query} from URL')
|
||||
|
||||
# Numeric query → single model fetch
|
||||
if query.isnumeric():
|
||||
model = client.get_model(int(query), token=token)
|
||||
@@ -49,7 +56,7 @@ def search_civitai(
|
||||
|
||||
all_models = response.items
|
||||
exact_models: list[CivitModel] = []
|
||||
if exact:
|
||||
if exact and query:
|
||||
q_lower = query.lower()
|
||||
for model in all_models:
|
||||
names = [model.name.lower()]
|
||||
|
||||
+93
-18
@@ -483,11 +483,21 @@ def create_ui():
|
||||
outputs=[models_outcome]
|
||||
)
|
||||
|
||||
with gr.Tab(label="CivitAI", elem_id="models_civitai_tab"):
|
||||
with gr.Tab(label="CivitAI", elem_id="models_civitai_tab") as civitai_tab:
|
||||
from modules.civitai.search_civitai import search_civitai, create_model_cards, base_models
|
||||
|
||||
def civitai_search(civit_search_text, civit_search_tag, civit_nsfw, civit_type, civit_base, civit_token):
|
||||
results = search_civitai(query=civit_search_text, tag=civit_search_tag, nsfw=civit_nsfw, types=civit_type, base=civit_base, token=civit_token)
|
||||
sort_fallback = ['', 'Most Downloaded', 'Highest Rated', 'Most Liked', 'Most Discussed',
|
||||
'Most Collected', 'Most Images', 'Newest', 'Oldest']
|
||||
type_fallback = ['', 'Checkpoint', 'TextualInversion', 'Hypernetwork', 'AestheticGradient',
|
||||
'LORA', 'LoCon', 'DoRA', 'Controlnet', 'Upscaler', 'MotionModule',
|
||||
'VAE', 'Poses', 'Wildcards', 'Workflows', 'Detection', 'Other']
|
||||
def civitai_search(civit_search_text, civit_search_tag, civit_nsfw, civit_type,
|
||||
civit_base, civit_token, civit_sort, civit_period):
|
||||
if civit_search_text and civit_search_tag:
|
||||
civit_search_tag = '' # query+tag is broken at API level, keyword wins
|
||||
results = search_civitai(query=civit_search_text, tag=civit_search_tag, nsfw=civit_nsfw,
|
||||
types=civit_type, base=civit_base, token=civit_token,
|
||||
sort=civit_sort, period=civit_period)
|
||||
html = create_model_cards(results)
|
||||
return html
|
||||
|
||||
@@ -496,50 +506,115 @@ def create_ui():
|
||||
opts.civitai_token = token
|
||||
opts.save()
|
||||
|
||||
def civitai_download(model_urls, model_names, model_types, model_path, civit_token, model_output):
|
||||
def civitai_download(model_urls, model_names, model_types, model_bases,
|
||||
model_ids, version_ids, model_path, civit_token, model_output):
|
||||
from modules.civitai.download_civitai import download_civit_model
|
||||
for model_url, model_name, model_type in zip(model_urls, model_names, model_types, strict=False):
|
||||
for model_url, model_name, model_type, model_base, model_id, version_id in zip(
|
||||
model_urls, model_names, model_types, model_bases, model_ids, version_ids, strict=False):
|
||||
msg = f"<h4>Initiating download</h4><div>{model_name} | {model_type} | <a href='{model_url}'>{model_url}</a></div><br>"
|
||||
yield msg + model_output
|
||||
download_civit_model(model_url, model_name, model_path, model_type, civit_token)
|
||||
download_civit_model(model_url, model_name, model_path, model_type, civit_token,
|
||||
base_model=model_base, model_id=int(model_id or 0), version_id=int(version_id or 0))
|
||||
yield model_output
|
||||
|
||||
def civitai_toggle_subfolder(enabled, template):
|
||||
opts.data['civitai_save_subfolder_enabled'] = enabled
|
||||
if enabled and not template:
|
||||
template = '{{BASEMODEL}}'
|
||||
opts.data['civitai_save_subfolder'] = template
|
||||
opts.save()
|
||||
return gr.update(value=template, interactive=enabled)
|
||||
|
||||
with gr.Row():
|
||||
gr.HTML('<h2>Search & Download</h2>')
|
||||
with gr.Row(elem_id='civitai_search_row'):
|
||||
civit_search_text = gr.Textbox(label='', placeholder='keyword', elem_id="civit_search_text")
|
||||
civit_search_tag = gr.Textbox(label='', placeholder='tag', elem_id="civit_search_text")
|
||||
civit_search_text = gr.Textbox(label='', placeholder='keyword, model id, or civitai url', elem_id="civit_search_text")
|
||||
civit_search_tag = gr.Textbox(label='', placeholder='tag', elem_id="civit_search_tag")
|
||||
civit_search_text_btn = ToolButton(value=ui_symbols.search, interactive=True, elem_id="civit_text_search")
|
||||
with gr.Accordion(label='Advanced', open=False, elem_id="civitai_search_options"):
|
||||
with gr.Accordion(label='Options', open=False, elem_id="civitai_search_options"):
|
||||
civit_download_btn = gr.Button(value="Download model", variant='primary', elem_id="civitai_download_btn", visible=False)
|
||||
with gr.Row():
|
||||
civit_token = gr.Textbox(opts.civitai_token, label='CivitAI token', placeholder='optional access token for private or gated models', elem_id="civitai_token")
|
||||
civit_type = gr.Dropdown(choices=type_fallback, label='Model type', value='', elem_id='civit_type')
|
||||
civit_base = gr.Dropdown(choices=base_models, label='Base model', value='')
|
||||
with gr.Row():
|
||||
civit_sort = gr.Dropdown(choices=sort_fallback, label='Sort', value='', elem_id='civit_sort')
|
||||
civit_period = gr.Dropdown(
|
||||
choices=['', 'AllTime', 'Year', 'Month', 'Week', 'Day'],
|
||||
label='Time period', value='', elem_id='civit_period',
|
||||
)
|
||||
with gr.Row():
|
||||
civit_nsfw = gr.Checkbox(label='NSFW allowed', value=True)
|
||||
with gr.Row():
|
||||
civit_type = gr.Textbox(label='Target model type', placeholder='Checkpoint, LORA, ...', value='')
|
||||
with gr.Row():
|
||||
# civit_base = gr.Textbox(label='Base model', placeholder='SDXL, ...')
|
||||
civit_base = gr.Dropdown(choices=base_models, label='Base model', value='')
|
||||
civit_token = gr.Textbox(opts.civitai_token, label='CivitAI token', placeholder='optional access token for private or gated models', elem_id="civitai_token")
|
||||
with gr.Row():
|
||||
civit_folder = gr.Textbox(label='Download folder', placeholder='optional folder for downloads')
|
||||
with gr.Row():
|
||||
civit_subfolder_enabled = gr.Checkbox(
|
||||
label='Sort downloads into subfolders',
|
||||
value=getattr(opts, 'civitai_save_subfolder_enabled', False),
|
||||
elem_id='civit_subfolder_enabled',
|
||||
)
|
||||
civit_subfolder_template = gr.Textbox(
|
||||
label='Subfolder template',
|
||||
value=getattr(opts, 'civitai_save_subfolder', '') if getattr(opts, 'civitai_save_subfolder_enabled', False) else '',
|
||||
placeholder='e.g. {{BASEMODEL}} or {{CREATOR}}/{{BASEMODEL}}',
|
||||
interactive=getattr(opts, 'civitai_save_subfolder_enabled', False),
|
||||
elem_id='civit_subfolder_template',
|
||||
)
|
||||
with gr.Row():
|
||||
civitai_models_output = gr.HTML('', elem_id="civitai_models_output")
|
||||
# sort, period, limit
|
||||
_dummy = gr.Label(visible=False) # dummy component to get argspec later
|
||||
civit_inputs = [civit_search_text, civit_search_tag, civit_nsfw, civit_type, civit_base, civit_token]
|
||||
_dummy = gr.Label(visible=False)
|
||||
civit_inputs = [civit_search_text, civit_search_tag, civit_nsfw, civit_type,
|
||||
civit_base, civit_token, civit_sort, civit_period]
|
||||
civit_search_text_btn.click(fn=civitai_search, inputs=civit_inputs, outputs=[civitai_models_output])
|
||||
civit_search_text.submit(fn=civitai_search, inputs=civit_inputs, outputs=[civitai_models_output])
|
||||
civit_search_tag.submit(fn=civitai_search, inputs=civit_inputs, outputs=[civitai_models_output])
|
||||
civit_token.change(fn=civitai_update_token, inputs=[civit_token], outputs=[])
|
||||
civit_subfolder_enabled.change(
|
||||
fn=civitai_toggle_subfolder,
|
||||
inputs=[civit_subfolder_enabled, civit_subfolder_template],
|
||||
outputs=[civit_subfolder_template],
|
||||
)
|
||||
civit_subfolder_template.change(
|
||||
fn=lambda v: (setattr(opts, 'civitai_save_subfolder', v), opts.save()),
|
||||
inputs=[civit_subfolder_template], outputs=[],
|
||||
)
|
||||
civit_download_btn.click(
|
||||
fn=civitai_download,
|
||||
_js="downloadCivitModel",
|
||||
inputs=[_dummy, _dummy, _dummy, civit_folder, civit_token, civitai_models_output],
|
||||
inputs=[_dummy, _dummy, _dummy, _dummy, _dummy, _dummy, civit_folder, civit_token, civitai_models_output],
|
||||
outputs=[civitai_models_output],
|
||||
show_progress='full',
|
||||
)
|
||||
|
||||
_civitai_loaded = False
|
||||
|
||||
def civitai_on_tab_enter():
|
||||
nonlocal _civitai_loaded
|
||||
if _civitai_loaded:
|
||||
return [gr.update(), gr.update(), gr.update(), gr.update(), gr.update()]
|
||||
_civitai_loaded = True
|
||||
from modules.civitai.client_civitai import client
|
||||
options = client.discover_options()
|
||||
type_choices = [''] + (options.get('types', []) or type_fallback[1:])
|
||||
sort_choices = [''] + (options.get('sort', []) or sort_fallback[1:])
|
||||
base_choices = [''] + (options.get('base_models', []) or base_models[1:])
|
||||
results = search_civitai(query='', sort='Most Downloaded', period='AllTime', limit=20)
|
||||
html = create_model_cards(results)
|
||||
return [
|
||||
gr.update(choices=type_choices),
|
||||
gr.update(choices=sort_choices, value='Most Downloaded'),
|
||||
gr.update(choices=base_choices),
|
||||
gr.update(value='AllTime'),
|
||||
html,
|
||||
]
|
||||
|
||||
civitai_tab.select(
|
||||
fn=civitai_on_tab_enter,
|
||||
inputs=[],
|
||||
outputs=[civit_type, civit_sort, civit_base, civit_period, civitai_models_output],
|
||||
)
|
||||
|
||||
with gr.Tab(label="Huggingface", elem_id="models_huggingface_tab"):
|
||||
from modules.models_hf import hf_search, hf_select, hf_download_model, hf_update_token
|
||||
with gr.Column(scale=6):
|
||||
|
||||
@@ -0,0 +1,653 @@
|
||||
#!/usr/bin/env python
|
||||
"""
|
||||
API tests for YOLO Detailer endpoints.
|
||||
|
||||
Tests:
|
||||
- GET /sdapi/v1/detailers — model enumeration
|
||||
- POST /sdapi/v1/detect — object detection on test images
|
||||
- POST /sdapi/v1/txt2img — generation with detailer enabled
|
||||
|
||||
Requires a running SD.Next instance with a model loaded.
|
||||
|
||||
Usage:
|
||||
python test/test-detailer-api.py [--url URL] [--image PATH]
|
||||
"""
|
||||
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import json
|
||||
import base64
|
||||
import argparse
|
||||
import requests
|
||||
import urllib3
|
||||
|
||||
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
|
||||
|
||||
# Reference model cover images with faces (best for detailer testing)
|
||||
FACE_TEST_IMAGES = [
|
||||
'models/Reference/ponyRealism_V23.jpg', # realistic woman, clear face
|
||||
'models/Reference/HiDream-ai--HiDream-I1-Fast.jpg', # realistic man, clear face + text
|
||||
'models/Reference/stabilityai--stable-diffusion-xl-base-1.0.jpg', # realistic woman portrait
|
||||
'models/Reference/CalamitousFelicitousness--Anima-sdnext-diffusers.jpg', # anime face (non-realistic test)
|
||||
]
|
||||
|
||||
# Fallback images (no guaranteed faces)
|
||||
FALLBACK_IMAGES = [
|
||||
'html/sdnext-robot-2k.jpg',
|
||||
'html/favicon.png',
|
||||
]
|
||||
|
||||
|
||||
class DetailerAPITest:
|
||||
"""Test harness for YOLO Detailer API endpoints."""
|
||||
|
||||
def __init__(self, base_url, image_path=None, timeout=300):
|
||||
self.base_url = base_url.rstrip('/')
|
||||
self.test_images = {} # name -> base64
|
||||
self.timeout = timeout
|
||||
self.results = {
|
||||
'enumerate': {'passed': 0, 'failed': 0, 'skipped': 0, 'tests': []},
|
||||
'detect': {'passed': 0, 'failed': 0, 'skipped': 0, 'tests': []},
|
||||
'generate': {'passed': 0, 'failed': 0, 'skipped': 0, 'tests': []},
|
||||
'detailer_params': {'passed': 0, 'failed': 0, 'skipped': 0, 'tests': []},
|
||||
}
|
||||
self._category = 'enumerate'
|
||||
self._critical_error = None
|
||||
self._load_images(image_path)
|
||||
|
||||
def _encode_image(self, path):
|
||||
from PIL import Image
|
||||
image = Image.open(path)
|
||||
if image.mode == 'RGBA':
|
||||
image = image.convert('RGB')
|
||||
buf = io.BytesIO()
|
||||
image.save(buf, 'JPEG')
|
||||
return base64.b64encode(buf.getvalue()).decode(), image.size
|
||||
|
||||
def _load_images(self, image_path=None):
|
||||
if image_path and os.path.exists(image_path):
|
||||
b64, size = self._encode_image(image_path)
|
||||
name = os.path.basename(image_path)
|
||||
self.test_images[name] = b64
|
||||
print(f" Test image: {image_path} ({size})")
|
||||
return
|
||||
|
||||
# Load all available face test images
|
||||
for p in FACE_TEST_IMAGES:
|
||||
if os.path.exists(p):
|
||||
b64, size = self._encode_image(p)
|
||||
name = os.path.basename(p)
|
||||
self.test_images[name] = b64
|
||||
print(f" Loaded: {name} ({size[0]}x{size[1]})")
|
||||
|
||||
# Fallback if no face images found
|
||||
if not self.test_images:
|
||||
for p in FALLBACK_IMAGES:
|
||||
if os.path.exists(p):
|
||||
b64, size = self._encode_image(p)
|
||||
name = os.path.basename(p)
|
||||
self.test_images[name] = b64
|
||||
print(f" Fallback: {name} ({size[0]}x{size[1]})")
|
||||
break
|
||||
|
||||
if not self.test_images:
|
||||
print(" WARNING: No test images found, detect tests will be skipped")
|
||||
|
||||
@property
|
||||
def image_b64(self):
|
||||
"""Return the first available test image for backwards compat."""
|
||||
if self.test_images:
|
||||
return next(iter(self.test_images.values()))
|
||||
return None
|
||||
|
||||
def _get(self, endpoint):
|
||||
try:
|
||||
r = requests.get(f'{self.base_url}{endpoint}', timeout=self.timeout, verify=False)
|
||||
if r.status_code != 200:
|
||||
return {'error': r.status_code, 'reason': r.reason}
|
||||
return r.json()
|
||||
except requests.exceptions.ConnectionError:
|
||||
return {'error': 'connection_refused', 'reason': 'Server not running'}
|
||||
except Exception as e:
|
||||
return {'error': 'exception', 'reason': str(e)}
|
||||
|
||||
def _post(self, endpoint, data):
|
||||
try:
|
||||
r = requests.post(f'{self.base_url}{endpoint}', json=data, timeout=self.timeout, verify=False)
|
||||
if r.status_code != 200:
|
||||
return {'error': r.status_code, 'reason': r.reason}
|
||||
return r.json()
|
||||
except requests.exceptions.ConnectionError:
|
||||
return {'error': 'connection_refused', 'reason': 'Server not running'}
|
||||
except Exception as e:
|
||||
return {'error': 'exception', 'reason': str(e)}
|
||||
|
||||
def record(self, passed, name, detail=''):
|
||||
status = 'PASS' if passed else 'FAIL'
|
||||
self.results[self._category]['passed' if passed else 'failed'] += 1
|
||||
self.results[self._category]['tests'].append((status, name))
|
||||
msg = f' {status}: {name}'
|
||||
if detail:
|
||||
msg += f' ({detail})'
|
||||
print(msg)
|
||||
|
||||
def skip(self, name, reason):
|
||||
self.results[self._category]['skipped'] += 1
|
||||
self.results[self._category]['tests'].append(('SKIP', name))
|
||||
print(f' SKIP: {name} ({reason})')
|
||||
|
||||
# =========================================================================
|
||||
# Tests: Model Enumeration
|
||||
# =========================================================================
|
||||
|
||||
def test_detailers_list(self):
|
||||
"""GET /sdapi/v1/detailers returns a list of available models."""
|
||||
self._category = 'enumerate'
|
||||
print("\n--- Detailer Model Enumeration ---")
|
||||
|
||||
data = self._get('/sdapi/v1/detailers')
|
||||
if 'error' in data:
|
||||
self.record(False, 'detailers_list', f"error: {data}")
|
||||
self._critical_error = f"Server error: {data}"
|
||||
return []
|
||||
|
||||
if not isinstance(data, list):
|
||||
self.record(False, 'detailers_list', f"expected list, got {type(data).__name__}")
|
||||
return []
|
||||
|
||||
self.record(True, 'detailers_list', f"{len(data)} models found")
|
||||
|
||||
# Verify each entry has expected fields
|
||||
if len(data) > 0:
|
||||
sample = data[0]
|
||||
has_name = 'name' in sample
|
||||
self.record(has_name, 'detailer_entry_has_name', f"sample: {sample}")
|
||||
if not has_name:
|
||||
self.record(False, 'detailer_entry_schema', "missing 'name' field")
|
||||
|
||||
return data
|
||||
|
||||
# =========================================================================
|
||||
# Tests: Detection
|
||||
# =========================================================================
|
||||
|
||||
def _validate_detect_response(self, data, label):
|
||||
"""Validate detection response schema and return detection count."""
|
||||
expected_keys = ['classes', 'labels', 'boxes', 'scores']
|
||||
for key in expected_keys:
|
||||
if key not in data:
|
||||
self.record(False, f'{label}_schema_{key}', f"missing '{key}'")
|
||||
return -1
|
||||
|
||||
# All arrays should have the same length
|
||||
lengths = [len(data[key]) for key in expected_keys]
|
||||
all_same = len(set(lengths)) <= 1
|
||||
if not all_same:
|
||||
self.record(False, f'{label}_array_lengths', f"mismatched: {dict(zip(expected_keys, lengths))}")
|
||||
return -1
|
||||
|
||||
n = lengths[0]
|
||||
|
||||
if n > 0:
|
||||
# Scores should be in [0, 1]
|
||||
scores_valid = all(0 <= s <= 1 for s in data['scores'])
|
||||
if not scores_valid:
|
||||
self.record(False, f'{label}_scores_range', f"scores: {data['scores']}")
|
||||
|
||||
# Boxes should be lists of 4 numbers
|
||||
boxes_valid = all(isinstance(b, list) and len(b) == 4 for b in data['boxes'])
|
||||
if not boxes_valid:
|
||||
self.record(False, f'{label}_boxes_format', "bad box format")
|
||||
|
||||
return n
|
||||
|
||||
# Face detection models to try (in priority order)
|
||||
FACE_MODELS = ['face-yolo8n', 'face-yolo8m', 'anzhc-face-1024-seg-8n']
|
||||
|
||||
def _pick_face_model(self, available_models):
|
||||
"""Pick the best face detection model from available ones."""
|
||||
available_names = [m.get('name', '') for m in available_models] if available_models else []
|
||||
for model in self.FACE_MODELS:
|
||||
if model in available_names:
|
||||
return model
|
||||
return '' # fall back to server default
|
||||
|
||||
def test_detect_all_images(self, available_models=None):
|
||||
"""POST /sdapi/v1/detect on each loaded test image with a face model."""
|
||||
self._category = 'detect'
|
||||
print("\n--- Detection Tests (per-image) ---")
|
||||
|
||||
if not self.test_images:
|
||||
self.skip('detect_all', 'no test images')
|
||||
return
|
||||
|
||||
if self._critical_error:
|
||||
self.skip('detect_all', self._critical_error)
|
||||
return
|
||||
|
||||
face_model = self._pick_face_model(available_models)
|
||||
if face_model:
|
||||
print(f" Using face model: {face_model}")
|
||||
else:
|
||||
print(" No face model available, using server default")
|
||||
|
||||
total_detections = 0
|
||||
any_face_found = False
|
||||
|
||||
for img_name, img_b64 in self.test_images.items():
|
||||
short = img_name.replace('.jpg', '')[:40]
|
||||
data = self._post('/sdapi/v1/detect', {'image': img_b64, 'model': face_model})
|
||||
|
||||
if 'error' in data:
|
||||
self.record(False, f'detect_{short}', f"error: {data}")
|
||||
continue
|
||||
|
||||
n = self._validate_detect_response(data, f'detect_{short}')
|
||||
if n < 0:
|
||||
continue
|
||||
|
||||
labels = data.get('labels', [])
|
||||
scores = data.get('scores', [])
|
||||
detail_parts = [f"{n} detections"]
|
||||
if labels:
|
||||
detail_parts.append(f"labels={labels}")
|
||||
if scores:
|
||||
detail_parts.append(f"top_score={max(scores):.3f}")
|
||||
|
||||
self.record(True, f'detect_{short}', ', '.join(detail_parts))
|
||||
total_detections += n
|
||||
if n > 0:
|
||||
any_face_found = True
|
||||
|
||||
self.record(any_face_found, 'detect_found_faces',
|
||||
f"{total_detections} total detections across {len(self.test_images)} images")
|
||||
|
||||
def test_detect_with_model(self, model_name):
|
||||
"""POST /sdapi/v1/detect with a specific model on all images."""
|
||||
if not self.test_images:
|
||||
self.skip(f'detect_model_{model_name}', 'no test images')
|
||||
return
|
||||
|
||||
total = 0
|
||||
for _img_name, img_b64 in self.test_images.items():
|
||||
data = self._post('/sdapi/v1/detect', {'image': img_b64, 'model': model_name})
|
||||
if 'error' not in data:
|
||||
total += len(data.get('scores', []))
|
||||
|
||||
self.record(True, f'detect_model_{model_name}', f"{total} detections across {len(self.test_images)} images")
|
||||
|
||||
# =========================================================================
|
||||
# Tests: Generation with Detailer
|
||||
# =========================================================================
|
||||
|
||||
def test_txt2img_with_detailer(self):
|
||||
"""POST /sdapi/v1/txt2img with detailer_enabled=True."""
|
||||
self._category = 'generate'
|
||||
print("\n--- Generation with Detailer ---")
|
||||
|
||||
if self._critical_error:
|
||||
self.skip('txt2img_detailer', self._critical_error)
|
||||
return
|
||||
|
||||
payload = {
|
||||
'prompt': 'a photo of a person, face, portrait',
|
||||
'negative_prompt': '',
|
||||
'steps': 10,
|
||||
'width': 512,
|
||||
'height': 512,
|
||||
'seed': 42,
|
||||
'save_images': False,
|
||||
'send_images': True,
|
||||
'detailer_enabled': True,
|
||||
'detailer_strength': 0.3,
|
||||
'detailer_steps': 5,
|
||||
'detailer_conf': 0.3,
|
||||
'detailer_max': 3,
|
||||
}
|
||||
|
||||
t0 = time.time()
|
||||
# Detailer generation is multi-pass (generate + detect + inpaint per region), use longer timeout
|
||||
try:
|
||||
r = requests.post(f'{self.base_url}/sdapi/v1/txt2img', json=payload, timeout=600, verify=False)
|
||||
if r.status_code != 200:
|
||||
data = {'error': r.status_code, 'reason': r.reason}
|
||||
else:
|
||||
data = r.json()
|
||||
except requests.exceptions.ConnectionError as e:
|
||||
self.record(False, 'txt2img_detailer', f"connection error (is a model loaded?): {e}")
|
||||
return
|
||||
except requests.exceptions.ReadTimeout:
|
||||
self.record(False, 'txt2img_detailer', 'timeout after 600s')
|
||||
return
|
||||
t1 = time.time()
|
||||
|
||||
if 'error' in data:
|
||||
self.record(False, 'txt2img_detailer', f"error: {data} (ensure a model is loaded)")
|
||||
return
|
||||
|
||||
# Should have images
|
||||
has_images = 'images' in data and len(data['images']) > 0
|
||||
self.record(has_images, 'txt2img_detailer_has_images', f"time={t1 - t0:.1f}s")
|
||||
|
||||
if has_images:
|
||||
# Decode and verify image
|
||||
from PIL import Image
|
||||
img_data = data['images'][0].split(',', 1)[0]
|
||||
img = Image.open(io.BytesIO(base64.b64decode(img_data)))
|
||||
self.record(True, 'txt2img_detailer_image_valid', f"size={img.size}")
|
||||
|
||||
# Check info field for detailer metadata
|
||||
if 'info' in data:
|
||||
info = data['info'] if isinstance(data['info'], str) else json.dumps(data['info'])
|
||||
has_detailer_info = 'detailer' in info.lower() or 'Detailer' in info
|
||||
self.record(has_detailer_info, 'txt2img_detailer_metadata',
|
||||
'detailer info found in metadata' if has_detailer_info else 'no detailer metadata (detection may have found nothing)')
|
||||
|
||||
def test_txt2img_without_detailer(self):
|
||||
"""POST /sdapi/v1/txt2img baseline without detailer (sanity check)."""
|
||||
if self._critical_error:
|
||||
self.skip('txt2img_baseline', self._critical_error)
|
||||
return
|
||||
|
||||
payload = {
|
||||
'prompt': 'a simple landscape',
|
||||
'steps': 5,
|
||||
'width': 512,
|
||||
'height': 512,
|
||||
'seed': 42,
|
||||
'save_images': False,
|
||||
'send_images': True,
|
||||
}
|
||||
|
||||
data = self._post('/sdapi/v1/txt2img', payload)
|
||||
if 'error' in data:
|
||||
self.record(False, 'txt2img_baseline', f"error: {data}")
|
||||
return
|
||||
|
||||
has_images = 'images' in data and len(data['images']) > 0
|
||||
self.record(has_images, 'txt2img_baseline', 'generation works without detailer')
|
||||
|
||||
# =========================================================================
|
||||
# Tests: Per-Request Detailer Param Validation
|
||||
# =========================================================================
|
||||
|
||||
def _txt2img(self, extra_params=None):
|
||||
"""Helper: generate a portrait with optional param overrides."""
|
||||
payload = {
|
||||
'prompt': 'a photo of a person, face, portrait',
|
||||
'steps': 10,
|
||||
'width': 512,
|
||||
'height': 512,
|
||||
'seed': 42,
|
||||
'save_images': False,
|
||||
'send_images': True,
|
||||
}
|
||||
if extra_params:
|
||||
payload.update(extra_params)
|
||||
try:
|
||||
r = requests.post(f'{self.base_url}/sdapi/v1/txt2img', json=payload, timeout=600, verify=False)
|
||||
if r.status_code != 200:
|
||||
return {'error': r.status_code, 'reason': r.reason}
|
||||
return r.json()
|
||||
except requests.exceptions.ConnectionError as e:
|
||||
return {'error': 'connection_refused', 'reason': str(e)}
|
||||
except requests.exceptions.ReadTimeout:
|
||||
return {'error': 'timeout', 'reason': 'timeout after 600s'}
|
||||
|
||||
def _decode_image(self, data):
|
||||
"""Decode first image from generation response into numpy array."""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
if 'images' not in data or len(data['images']) == 0:
|
||||
return None
|
||||
img_data = data['images'][0].split(',', 1)[0]
|
||||
img = Image.open(io.BytesIO(base64.b64decode(img_data))).convert('RGB')
|
||||
return np.array(img, dtype=np.float32)
|
||||
|
||||
def _pixel_diff(self, arr_a, arr_b):
|
||||
"""Mean absolute pixel difference between two images."""
|
||||
import numpy as np
|
||||
if arr_a is None or arr_b is None or arr_a.shape != arr_b.shape:
|
||||
return -1.0
|
||||
return float(np.abs(arr_a - arr_b).mean())
|
||||
|
||||
def _get_info(self, data):
|
||||
"""Extract info string from generation response."""
|
||||
if 'info' not in data:
|
||||
return ''
|
||||
info = data['info']
|
||||
return info if isinstance(info, str) else json.dumps(info)
|
||||
|
||||
def run_detailer_param_tests(self, available_models=None):
|
||||
"""Verify per-request detailer params change the output."""
|
||||
self._category = 'detailer_params'
|
||||
print("\n--- Per-Request Detailer Param Validation ---")
|
||||
|
||||
if self._critical_error:
|
||||
self.skip('detailer_params_all', self._critical_error)
|
||||
return
|
||||
|
||||
# Generate baseline WITHOUT detailer (same seed/prompt as detailer tests)
|
||||
print(" Generating baseline (no detailer)...")
|
||||
baseline_data = self._txt2img()
|
||||
if 'error' in baseline_data:
|
||||
self.record(False, 'detailer_baseline', f"error: {baseline_data}")
|
||||
return
|
||||
baseline = self._decode_image(baseline_data)
|
||||
if baseline is None:
|
||||
self.record(False, 'detailer_baseline', 'no image')
|
||||
return
|
||||
self.record(True, 'detailer_baseline')
|
||||
|
||||
# Generate WITH detailer enabled (default params)
|
||||
print(" Generating with detailer (defaults)...")
|
||||
detailer_default_data = self._txt2img({
|
||||
'detailer_enabled': True,
|
||||
'detailer_strength': 0.3,
|
||||
'detailer_steps': 5,
|
||||
'detailer_conf': 0.3,
|
||||
})
|
||||
if 'error' in detailer_default_data:
|
||||
self.record(False, 'detailer_default', f"error: {detailer_default_data}")
|
||||
return
|
||||
detailer_default = self._decode_image(detailer_default_data)
|
||||
|
||||
# Detailer ON vs OFF should produce different images (if a face was detected)
|
||||
diff_on_off = self._pixel_diff(baseline, detailer_default)
|
||||
self.record(diff_on_off > 0.5, 'detailer_on_vs_off',
|
||||
f"mean_diff={diff_on_off:.2f}" if diff_on_off > 0.5
|
||||
else f"identical (diff={diff_on_off:.4f}) — no face detected?")
|
||||
|
||||
# -- Strength variation --
|
||||
print(" Testing strength variation...")
|
||||
strong_data = self._txt2img({
|
||||
'detailer_enabled': True,
|
||||
'detailer_strength': 0.7,
|
||||
'detailer_steps': 5,
|
||||
'detailer_conf': 0.3,
|
||||
})
|
||||
if 'error' not in strong_data:
|
||||
strong = self._decode_image(strong_data)
|
||||
diff_strong = self._pixel_diff(detailer_default, strong)
|
||||
self.record(diff_strong > 0.5, 'detailer_strength_effect',
|
||||
f"strength 0.3 vs 0.7: diff={diff_strong:.2f}")
|
||||
|
||||
# -- Steps variation --
|
||||
print(" Testing steps variation...")
|
||||
more_steps_data = self._txt2img({
|
||||
'detailer_enabled': True,
|
||||
'detailer_strength': 0.3,
|
||||
'detailer_steps': 20,
|
||||
'detailer_conf': 0.3,
|
||||
})
|
||||
if 'error' not in more_steps_data:
|
||||
more_steps = self._decode_image(more_steps_data)
|
||||
diff_steps = self._pixel_diff(detailer_default, more_steps)
|
||||
self.record(diff_steps > 0.5, 'detailer_steps_effect',
|
||||
f"steps 5 vs 20: diff={diff_steps:.2f}")
|
||||
|
||||
# -- Resolution variation --
|
||||
print(" Testing resolution variation...")
|
||||
hires_data = self._txt2img({
|
||||
'detailer_enabled': True,
|
||||
'detailer_strength': 0.3,
|
||||
'detailer_steps': 5,
|
||||
'detailer_conf': 0.3,
|
||||
'detailer_resolution': 512,
|
||||
})
|
||||
if 'error' not in hires_data:
|
||||
hires = self._decode_image(hires_data)
|
||||
diff_res = self._pixel_diff(detailer_default, hires)
|
||||
self.record(diff_res > 0.5, 'detailer_resolution_effect',
|
||||
f"resolution 1024 vs 512: diff={diff_res:.2f}")
|
||||
|
||||
# -- Segmentation mode --
|
||||
# Segmentation requires a -seg model (e.g. anzhc-face-1024-seg-8n).
|
||||
# Detection-only models (face-yolo8n) don't produce masks, so the flag has no effect.
|
||||
seg_models = [m.get('name', '') for m in (available_models or [])
|
||||
if 'seg' in m.get('name', '').lower() and 'face' in m.get('name', '').lower()]
|
||||
if seg_models:
|
||||
seg_model = seg_models[0]
|
||||
print(f" Testing segmentation mode (model={seg_model})...")
|
||||
# bbox baseline with the seg model
|
||||
seg_bbox_data = self._txt2img({
|
||||
'detailer_enabled': True,
|
||||
'detailer_strength': 0.3,
|
||||
'detailer_steps': 5,
|
||||
'detailer_conf': 0.3,
|
||||
'detailer_segmentation': False,
|
||||
'detailer_models': [seg_model],
|
||||
})
|
||||
seg_data = self._txt2img({
|
||||
'detailer_enabled': True,
|
||||
'detailer_strength': 0.3,
|
||||
'detailer_steps': 5,
|
||||
'detailer_conf': 0.3,
|
||||
'detailer_segmentation': True,
|
||||
'detailer_models': [seg_model],
|
||||
})
|
||||
if 'error' not in seg_data and 'error' not in seg_bbox_data:
|
||||
seg_bbox = self._decode_image(seg_bbox_data)
|
||||
seg_mask = self._decode_image(seg_data)
|
||||
diff_seg = self._pixel_diff(seg_bbox, seg_mask)
|
||||
self.record(diff_seg > 0.5, 'detailer_segmentation_effect',
|
||||
f"bbox vs seg mask ({seg_model}): diff={diff_seg:.2f}")
|
||||
else:
|
||||
err = seg_data if 'error' in seg_data else seg_bbox_data
|
||||
self.record(False, 'detailer_segmentation_effect', f"error: {err}")
|
||||
else:
|
||||
print(" Testing segmentation mode...")
|
||||
seg_data = {'error': 'skipped'}
|
||||
self.skip('detailer_segmentation_effect', 'no face-seg model available')
|
||||
|
||||
# -- Confidence threshold --
|
||||
print(" Testing confidence threshold...")
|
||||
high_conf_data = self._txt2img({
|
||||
'detailer_enabled': True,
|
||||
'detailer_strength': 0.3,
|
||||
'detailer_steps': 5,
|
||||
'detailer_conf': 0.95,
|
||||
})
|
||||
if 'error' not in high_conf_data:
|
||||
high_conf = self._decode_image(high_conf_data)
|
||||
diff_conf = self._pixel_diff(baseline, high_conf)
|
||||
# High confidence may reject detections, making output closer to baseline
|
||||
self.record(True, 'detailer_conf_effect',
|
||||
f"conf=0.95 vs baseline: diff={diff_conf:.2f} "
|
||||
f"(low diff = detections filtered out, high diff = still detected)")
|
||||
|
||||
# -- Custom detailer prompt --
|
||||
print(" Testing detailer prompt override...")
|
||||
prompt_data = self._txt2img({
|
||||
'detailer_enabled': True,
|
||||
'detailer_strength': 0.5,
|
||||
'detailer_steps': 5,
|
||||
'detailer_conf': 0.3,
|
||||
'detailer_prompt': 'a detailed close-up face with freckles',
|
||||
})
|
||||
if 'error' not in prompt_data:
|
||||
prompt_result = self._decode_image(prompt_data)
|
||||
diff_prompt = self._pixel_diff(detailer_default, prompt_result)
|
||||
self.record(diff_prompt > 0.5, 'detailer_prompt_effect',
|
||||
f"custom prompt vs default: diff={diff_prompt:.2f}")
|
||||
|
||||
# -- Metadata verification across params --
|
||||
for test_data, label in [
|
||||
(detailer_default_data, 'detailer_default'),
|
||||
(strong_data if 'error' not in strong_data else None, 'detailer_strong'),
|
||||
(more_steps_data if 'error' not in more_steps_data else None, 'detailer_more_steps'),
|
||||
(seg_data if 'error' not in seg_data else None, 'detailer_segmentation'),
|
||||
]:
|
||||
if test_data is None:
|
||||
continue
|
||||
info = self._get_info(test_data)
|
||||
has_meta = 'detailer' in info.lower() or 'Detailer' in info
|
||||
self.record(has_meta, f'{label}_metadata',
|
||||
'detailer info in metadata' if has_meta else 'no detailer metadata')
|
||||
|
||||
# -- Param isolation: generate without detailer after all detailer runs --
|
||||
print(" Testing param isolation...")
|
||||
after_data = self._txt2img()
|
||||
if 'error' not in after_data:
|
||||
after = self._decode_image(after_data)
|
||||
leak_diff = self._pixel_diff(baseline, after)
|
||||
self.record(leak_diff < 0.5, 'detailer_param_isolation',
|
||||
f"post-detailer baseline diff={leak_diff:.4f}" if leak_diff < 0.5
|
||||
else f"LEAK: baseline changed (diff={leak_diff:.2f})")
|
||||
|
||||
# =========================================================================
|
||||
# Runner
|
||||
# =========================================================================
|
||||
|
||||
def run_all(self):
|
||||
print("=" * 60)
|
||||
print("YOLO Detailer API Test Suite")
|
||||
print(f"Server: {self.base_url}")
|
||||
print("=" * 60)
|
||||
|
||||
# Enumerate
|
||||
models = self.test_detailers_list()
|
||||
|
||||
# Detect across all loaded test images
|
||||
self.test_detect_all_images(models)
|
||||
# Test with first available model if any
|
||||
if models and len(models) > 0:
|
||||
model_name = models[0].get('name', models[0].get('filename', ''))
|
||||
if model_name:
|
||||
self.test_detect_with_model(model_name)
|
||||
|
||||
# Generate
|
||||
self.test_txt2img_without_detailer()
|
||||
self.test_txt2img_with_detailer()
|
||||
|
||||
# Per-request detailer param validation
|
||||
self.run_detailer_param_tests(models)
|
||||
|
||||
# Summary
|
||||
print("\n" + "=" * 60)
|
||||
print("Results")
|
||||
print("=" * 60)
|
||||
total_passed = 0
|
||||
total_failed = 0
|
||||
total_skipped = 0
|
||||
for cat, data in self.results.items():
|
||||
total_passed += data['passed']
|
||||
total_failed += data['failed']
|
||||
total_skipped += data['skipped']
|
||||
status = 'PASS' if data['failed'] == 0 else 'FAIL'
|
||||
print(f" {cat}: {data['passed']} passed, {data['failed']} failed, {data['skipped']} skipped [{status}]")
|
||||
print(f" Total: {total_passed} passed, {total_failed} failed, {total_skipped} skipped")
|
||||
print("=" * 60)
|
||||
return total_failed == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description='YOLO Detailer API Tests')
|
||||
parser.add_argument('--url', default=os.environ.get('SDAPI_URL', 'http://127.0.0.1:7860'), help='server URL')
|
||||
parser.add_argument('--image', default=None, help='test image path')
|
||||
args = parser.parse_args()
|
||||
test = DetailerAPITest(args.url, args.image)
|
||||
success = test.run_all()
|
||||
sys.exit(0 if success else 1)
|
||||
@@ -0,0 +1,615 @@
|
||||
#!/usr/bin/env python
|
||||
"""
|
||||
API tests for generation with scheduler params, color grading, and latent corrections.
|
||||
|
||||
Tests:
|
||||
- GET /sdapi/v1/samplers — sampler enumeration and config
|
||||
- POST /sdapi/v1/txt2img — generation with various samplers
|
||||
- POST /sdapi/v1/txt2img — generation with color grading params
|
||||
- POST /sdapi/v1/txt2img — generation with latent correction params
|
||||
|
||||
Requires a running SD.Next instance with a model loaded.
|
||||
|
||||
Usage:
|
||||
python test/test-generation-api.py [--url URL] [--steps STEPS]
|
||||
"""
|
||||
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import time
|
||||
import base64
|
||||
import argparse
|
||||
import requests
|
||||
import urllib3
|
||||
|
||||
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
|
||||
|
||||
|
||||
class GenerationAPITest:
|
||||
"""Test harness for generation API with scheduler and grading params."""
|
||||
|
||||
# Samplers to test — a representative subset covering different scheduler families
|
||||
TEST_SAMPLERS = [
|
||||
'Euler a',
|
||||
'Euler',
|
||||
'DPM++ 2M',
|
||||
'UniPC',
|
||||
'DDIM',
|
||||
'DPM++ 2M SDE',
|
||||
]
|
||||
|
||||
def __init__(self, base_url, steps=10, timeout=300):
|
||||
self.base_url = base_url.rstrip('/')
|
||||
self.steps = steps
|
||||
self.timeout = timeout
|
||||
self.results = {
|
||||
'samplers': {'passed': 0, 'failed': 0, 'skipped': 0, 'tests': []},
|
||||
'generation': {'passed': 0, 'failed': 0, 'skipped': 0, 'tests': []},
|
||||
'grading': {'passed': 0, 'failed': 0, 'skipped': 0, 'tests': []},
|
||||
'correction': {'passed': 0, 'failed': 0, 'skipped': 0, 'tests': []},
|
||||
'param_validation': {'passed': 0, 'failed': 0, 'skipped': 0, 'tests': []},
|
||||
}
|
||||
self._category = 'samplers'
|
||||
self._critical_error = None
|
||||
|
||||
def _get(self, endpoint):
|
||||
try:
|
||||
r = requests.get(f'{self.base_url}{endpoint}', timeout=self.timeout, verify=False)
|
||||
if r.status_code != 200:
|
||||
return {'error': r.status_code, 'reason': r.reason}
|
||||
return r.json()
|
||||
except requests.exceptions.ConnectionError:
|
||||
return {'error': 'connection_refused', 'reason': 'Server not running'}
|
||||
except Exception as e:
|
||||
return {'error': 'exception', 'reason': str(e)}
|
||||
|
||||
def _post(self, endpoint, data):
|
||||
try:
|
||||
r = requests.post(f'{self.base_url}{endpoint}', json=data, timeout=self.timeout, verify=False)
|
||||
if r.status_code != 200:
|
||||
return {'error': r.status_code, 'reason': r.reason}
|
||||
return r.json()
|
||||
except requests.exceptions.ConnectionError:
|
||||
return {'error': 'connection_refused', 'reason': 'Server not running'}
|
||||
except Exception as e:
|
||||
return {'error': 'exception', 'reason': str(e)}
|
||||
|
||||
def record(self, passed, name, detail=''):
|
||||
status = 'PASS' if passed else 'FAIL'
|
||||
self.results[self._category]['passed' if passed else 'failed'] += 1
|
||||
self.results[self._category]['tests'].append((status, name))
|
||||
msg = f' {status}: {name}'
|
||||
if detail:
|
||||
msg += f' ({detail})'
|
||||
print(msg)
|
||||
|
||||
def skip(self, name, reason):
|
||||
self.results[self._category]['skipped'] += 1
|
||||
self.results[self._category]['tests'].append(('SKIP', name))
|
||||
print(f' SKIP: {name} ({reason})')
|
||||
|
||||
def _txt2img(self, extra_params=None, prompt='a cat'):
|
||||
"""Helper: run txt2img with base params + overrides. Returns (data, time)."""
|
||||
payload = {
|
||||
'prompt': prompt,
|
||||
'steps': self.steps,
|
||||
'width': 512,
|
||||
'height': 512,
|
||||
'seed': 42,
|
||||
'save_images': False,
|
||||
'send_images': True,
|
||||
}
|
||||
if extra_params:
|
||||
payload.update(extra_params)
|
||||
t0 = time.time()
|
||||
data = self._post('/sdapi/v1/txt2img', payload)
|
||||
return data, time.time() - t0
|
||||
|
||||
def _check_generation(self, data, test_name, elapsed):
|
||||
"""Validate a generation response has images."""
|
||||
if 'error' in data:
|
||||
self.record(False, test_name, f"error: {data}")
|
||||
return False
|
||||
has_images = 'images' in data and len(data['images']) > 0
|
||||
self.record(has_images, test_name, f"time={elapsed:.1f}s")
|
||||
return has_images
|
||||
|
||||
def _get_info(self, data):
|
||||
"""Extract info string from generation response."""
|
||||
if 'info' not in data:
|
||||
return ''
|
||||
info = data['info']
|
||||
return info if isinstance(info, str) else json.dumps(info)
|
||||
|
||||
def _decode_image(self, data):
|
||||
"""Decode first image from generation response into numpy array."""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
if 'images' not in data or len(data['images']) == 0:
|
||||
return None
|
||||
img_data = data['images'][0].split(',', 1)[0]
|
||||
img = Image.open(io.BytesIO(base64.b64decode(img_data))).convert('RGB')
|
||||
return np.array(img, dtype=np.float32)
|
||||
|
||||
def _pixel_diff(self, arr_a, arr_b):
|
||||
"""Mean absolute pixel difference between two images (0-255 scale)."""
|
||||
import numpy as np
|
||||
if arr_a is None or arr_b is None:
|
||||
return -1.0
|
||||
if arr_a.shape != arr_b.shape:
|
||||
return -1.0
|
||||
return float(np.abs(arr_a - arr_b).mean())
|
||||
|
||||
def _channel_means(self, arr):
|
||||
"""Return per-channel means [R, G, B]."""
|
||||
if arr is None:
|
||||
return [0, 0, 0]
|
||||
return [float(arr[:, :, c].mean()) for c in range(3)]
|
||||
|
||||
# =========================================================================
|
||||
# Tests: Sampler Enumeration
|
||||
# =========================================================================
|
||||
|
||||
def test_samplers_list(self):
|
||||
"""GET /sdapi/v1/samplers returns available samplers with config."""
|
||||
self._category = 'samplers'
|
||||
print("\n--- Sampler Enumeration ---")
|
||||
|
||||
data = self._get('/sdapi/v1/samplers')
|
||||
if 'error' in data:
|
||||
self.record(False, 'samplers_list', f"error: {data}")
|
||||
self._critical_error = f"Server error: {data}"
|
||||
return []
|
||||
|
||||
if not isinstance(data, list):
|
||||
self.record(False, 'samplers_list', f"expected list, got {type(data).__name__}")
|
||||
return []
|
||||
|
||||
self.record(True, 'samplers_list', f"{len(data)} samplers available")
|
||||
|
||||
# Check that each sampler has a name
|
||||
sampler_names = []
|
||||
for s in data:
|
||||
name = s.get('name', '')
|
||||
if name:
|
||||
sampler_names.append(name)
|
||||
|
||||
self.record(len(sampler_names) == len(data), 'samplers_have_names',
|
||||
f"{len(sampler_names)}/{len(data)} have names")
|
||||
|
||||
# Check for our test samplers
|
||||
for test_sampler in self.TEST_SAMPLERS:
|
||||
found = test_sampler in sampler_names
|
||||
if not found:
|
||||
self.skip(f'sampler_available_{test_sampler}', 'not in server sampler list')
|
||||
else:
|
||||
self.record(True, f'sampler_available_{test_sampler}')
|
||||
|
||||
return sampler_names
|
||||
|
||||
# =========================================================================
|
||||
# Tests: Generation with Different Samplers
|
||||
# =========================================================================
|
||||
|
||||
def test_samplers_generate(self, available_samplers):
|
||||
"""Generate with each test sampler and verify success."""
|
||||
self._category = 'generation'
|
||||
print("\n--- Generation with Different Samplers ---")
|
||||
|
||||
if self._critical_error:
|
||||
for s in self.TEST_SAMPLERS:
|
||||
self.skip(f'generate_{s}', self._critical_error)
|
||||
return
|
||||
|
||||
for sampler in self.TEST_SAMPLERS:
|
||||
if sampler not in available_samplers:
|
||||
self.skip(f'generate_{sampler}', 'sampler not available')
|
||||
continue
|
||||
data, elapsed = self._txt2img({'sampler_name': sampler})
|
||||
self._check_generation(data, f'generate_{sampler}', elapsed)
|
||||
|
||||
# =========================================================================
|
||||
# Tests: Color Grading Params
|
||||
# =========================================================================
|
||||
|
||||
def test_grading_brightness_contrast(self):
|
||||
"""Generate with grading brightness and contrast."""
|
||||
data, elapsed = self._txt2img({
|
||||
'grading_brightness': 0.2,
|
||||
'grading_contrast': 0.3,
|
||||
})
|
||||
self._check_generation(data, 'grading_brightness_contrast', elapsed)
|
||||
|
||||
def test_grading_saturation_hue(self):
|
||||
"""Generate with grading saturation and hue shift."""
|
||||
data, elapsed = self._txt2img({
|
||||
'grading_saturation': 0.5,
|
||||
'grading_hue': 0.1,
|
||||
})
|
||||
self._check_generation(data, 'grading_saturation_hue', elapsed)
|
||||
|
||||
def test_grading_gamma_sharpness(self):
|
||||
"""Generate with gamma correction and sharpness."""
|
||||
data, elapsed = self._txt2img({
|
||||
'grading_gamma': 0.8,
|
||||
'grading_sharpness': 0.5,
|
||||
})
|
||||
self._check_generation(data, 'grading_gamma_sharpness', elapsed)
|
||||
|
||||
def test_grading_color_temp(self):
|
||||
"""Generate with warm color temperature."""
|
||||
data, elapsed = self._txt2img({
|
||||
'grading_color_temp': 3500,
|
||||
})
|
||||
self._check_generation(data, 'grading_color_temp', elapsed)
|
||||
|
||||
def test_grading_tone(self):
|
||||
"""Generate with shadows/midtones/highlights adjustments."""
|
||||
data, elapsed = self._txt2img({
|
||||
'grading_shadows': 0.3,
|
||||
'grading_midtones': -0.1,
|
||||
'grading_highlights': 0.2,
|
||||
})
|
||||
self._check_generation(data, 'grading_tone', elapsed)
|
||||
|
||||
def test_grading_effects(self):
|
||||
"""Generate with vignette and grain."""
|
||||
data, elapsed = self._txt2img({
|
||||
'grading_vignette': 0.5,
|
||||
'grading_grain': 0.3,
|
||||
})
|
||||
self._check_generation(data, 'grading_effects', elapsed)
|
||||
|
||||
def test_grading_split_toning(self):
|
||||
"""Generate with split toning colors."""
|
||||
data, elapsed = self._txt2img({
|
||||
'grading_shadows_tint': '#003366',
|
||||
'grading_highlights_tint': '#ffcc00',
|
||||
'grading_split_tone_balance': 0.6,
|
||||
})
|
||||
self._check_generation(data, 'grading_split_toning', elapsed)
|
||||
|
||||
def test_grading_combined(self):
|
||||
"""Generate with multiple grading params at once."""
|
||||
data, elapsed = self._txt2img({
|
||||
'grading_brightness': 0.1,
|
||||
'grading_contrast': 0.2,
|
||||
'grading_saturation': 0.3,
|
||||
'grading_gamma': 0.9,
|
||||
'grading_color_temp': 5000,
|
||||
'grading_vignette': 0.3,
|
||||
})
|
||||
self._check_generation(data, 'grading_combined', elapsed)
|
||||
|
||||
def run_grading_tests(self):
|
||||
"""Run all grading tests."""
|
||||
self._category = 'grading'
|
||||
print("\n--- Color Grading Tests ---")
|
||||
|
||||
if self._critical_error:
|
||||
self.skip('grading_all', self._critical_error)
|
||||
return
|
||||
|
||||
self.test_grading_brightness_contrast()
|
||||
self.test_grading_saturation_hue()
|
||||
self.test_grading_gamma_sharpness()
|
||||
self.test_grading_color_temp()
|
||||
self.test_grading_tone()
|
||||
self.test_grading_effects()
|
||||
self.test_grading_split_toning()
|
||||
self.test_grading_combined()
|
||||
|
||||
# =========================================================================
|
||||
# Tests: Latent Correction Params
|
||||
# =========================================================================
|
||||
|
||||
def test_correction_brightness(self):
|
||||
"""Generate with latent brightness correction."""
|
||||
data, elapsed = self._txt2img({'hdr_brightness': 1.5})
|
||||
ok = self._check_generation(data, 'correction_brightness', elapsed)
|
||||
if ok:
|
||||
info = self._get_info(data)
|
||||
has_param = 'Latent brightness' in info
|
||||
self.record(has_param, 'correction_brightness_metadata',
|
||||
'found in info' if has_param else 'not found in info')
|
||||
|
||||
def test_correction_color(self):
|
||||
"""Generate with latent color centering."""
|
||||
data, elapsed = self._txt2img({'hdr_color': 0.5, 'hdr_mode': 1})
|
||||
ok = self._check_generation(data, 'correction_color', elapsed)
|
||||
if ok:
|
||||
info = self._get_info(data)
|
||||
has_param = 'Latent color' in info
|
||||
self.record(has_param, 'correction_color_metadata',
|
||||
'found in info' if has_param else 'not found in info')
|
||||
|
||||
def test_correction_clamp(self):
|
||||
"""Generate with latent clamping."""
|
||||
data, elapsed = self._txt2img({
|
||||
'hdr_clamp': True,
|
||||
'hdr_threshold': 0.8,
|
||||
'hdr_boundary': 4.0,
|
||||
})
|
||||
ok = self._check_generation(data, 'correction_clamp', elapsed)
|
||||
if ok:
|
||||
info = self._get_info(data)
|
||||
has_param = 'Latent clamp' in info
|
||||
self.record(has_param, 'correction_clamp_metadata',
|
||||
'found in info' if has_param else 'not found in info')
|
||||
|
||||
def test_correction_sharpen(self):
|
||||
"""Generate with latent sharpening."""
|
||||
data, elapsed = self._txt2img({'hdr_sharpen': 1.0})
|
||||
ok = self._check_generation(data, 'correction_sharpen', elapsed)
|
||||
if ok:
|
||||
info = self._get_info(data)
|
||||
has_param = 'Latent sharpen' in info
|
||||
self.record(has_param, 'correction_sharpen_metadata',
|
||||
'found in info' if has_param else 'not found in info')
|
||||
|
||||
def test_correction_maximize(self):
|
||||
"""Generate with latent maximize/normalize."""
|
||||
data, elapsed = self._txt2img({
|
||||
'hdr_maximize': True,
|
||||
'hdr_max_center': 0.6,
|
||||
'hdr_max_boundary': 2.0,
|
||||
})
|
||||
ok = self._check_generation(data, 'correction_maximize', elapsed)
|
||||
if ok:
|
||||
info = self._get_info(data)
|
||||
has_param = 'Latent max' in info
|
||||
self.record(has_param, 'correction_maximize_metadata',
|
||||
'found in info' if has_param else 'not found in info')
|
||||
|
||||
def test_correction_combined(self):
|
||||
"""Generate with multiple correction params."""
|
||||
data, elapsed = self._txt2img({
|
||||
'hdr_brightness': 1.0,
|
||||
'hdr_color': 0.3,
|
||||
'hdr_sharpen': 0.5,
|
||||
'hdr_clamp': True,
|
||||
})
|
||||
ok = self._check_generation(data, 'correction_combined', elapsed)
|
||||
if ok:
|
||||
info = self._get_info(data)
|
||||
# At least some correction params should appear
|
||||
found = [k for k in ['Latent brightness', 'Latent color', 'Latent sharpen', 'Latent clamp'] if k in info]
|
||||
self.record(len(found) > 0, 'correction_combined_metadata', f"found: {found}")
|
||||
|
||||
def run_correction_tests(self):
|
||||
"""Run all latent correction tests."""
|
||||
self._category = 'correction'
|
||||
print("\n--- Latent Correction Tests ---")
|
||||
|
||||
if self._critical_error:
|
||||
self.skip('correction_all', self._critical_error)
|
||||
return
|
||||
|
||||
self.test_correction_brightness()
|
||||
self.test_correction_color()
|
||||
self.test_correction_clamp()
|
||||
self.test_correction_sharpen()
|
||||
self.test_correction_maximize()
|
||||
self.test_correction_combined()
|
||||
|
||||
# =========================================================================
|
||||
# Tests: Per-Request Param Validation (baseline comparison)
|
||||
# =========================================================================
|
||||
|
||||
def _generate_baseline(self):
|
||||
"""Generate a baseline image with no grading/correction. Cache and reuse."""
|
||||
if hasattr(self, '_baseline_arr') and self._baseline_arr is not None:
|
||||
return self._baseline_arr, self._baseline_data
|
||||
data, elapsed = self._txt2img()
|
||||
if 'error' in data or 'images' not in data:
|
||||
return None, data
|
||||
self._baseline_arr = self._decode_image(data)
|
||||
self._baseline_data = data
|
||||
print(f' Baseline generated: time={elapsed:.1f}s mean={self._channel_means(self._baseline_arr)}')
|
||||
return self._baseline_arr, data
|
||||
|
||||
def _compare_param(self, name, params, check_fn=None):
|
||||
"""Generate with params and compare to baseline. Optionally run check_fn(baseline, result)."""
|
||||
baseline, _ = self._generate_baseline()
|
||||
if baseline is None:
|
||||
self.skip(f'param_{name}', 'baseline generation failed')
|
||||
return
|
||||
|
||||
data, elapsed = self._txt2img(params)
|
||||
if 'error' in data:
|
||||
self.record(False, f'param_{name}', f"generation error: {data}")
|
||||
return
|
||||
|
||||
result = self._decode_image(data)
|
||||
if result is None:
|
||||
self.record(False, f'param_{name}', 'no image in response')
|
||||
return
|
||||
|
||||
diff = self._pixel_diff(baseline, result)
|
||||
differs = diff > 0.5 # more than 0.5/255 mean difference
|
||||
self.record(differs, f'param_{name}_differs',
|
||||
f"mean_diff={diff:.2f}" if differs else f"images identical (diff={diff:.4f})")
|
||||
|
||||
if check_fn and differs:
|
||||
try:
|
||||
ok, detail = check_fn(baseline, result, data)
|
||||
self.record(ok, f'param_{name}_direction', detail)
|
||||
except Exception as e:
|
||||
self.record(False, f'param_{name}_direction', f"check error: {e}")
|
||||
|
||||
def run_param_validation_tests(self):
|
||||
"""Verify per-request grading/correction params actually change the output."""
|
||||
self._category = 'param_validation'
|
||||
print("\n--- Per-Request Param Validation ---")
|
||||
|
||||
if self._critical_error:
|
||||
self.skip('param_validation_all', self._critical_error)
|
||||
return
|
||||
|
||||
import numpy as np
|
||||
|
||||
# -- Grading params --
|
||||
|
||||
# Brightness: positive should increase mean pixel value
|
||||
def check_brightness(base, result, _data):
|
||||
base_mean = float(base.mean())
|
||||
result_mean = float(result.mean())
|
||||
return result_mean > base_mean, f"baseline={base_mean:.1f} graded={result_mean:.1f}"
|
||||
self._compare_param('grading_brightness', {'grading_brightness': 0.3}, check_brightness)
|
||||
|
||||
# Contrast: should increase standard deviation
|
||||
def check_contrast(base, result, _data):
|
||||
return float(result.std()) > float(base.std()), \
|
||||
f"baseline_std={float(base.std()):.1f} graded_std={float(result.std()):.1f}"
|
||||
self._compare_param('grading_contrast', {'grading_contrast': 0.5}, check_contrast)
|
||||
|
||||
# Saturation: desaturation should reduce color channel spread
|
||||
def check_desaturation(base, result, _data):
|
||||
base_spread = max(self._channel_means(base)) - min(self._channel_means(base))
|
||||
result_spread = max(self._channel_means(result)) - min(self._channel_means(result))
|
||||
return result_spread < base_spread, \
|
||||
f"baseline_spread={base_spread:.1f} graded_spread={result_spread:.1f}"
|
||||
self._compare_param('grading_saturation_neg', {'grading_saturation': -0.5}, check_desaturation)
|
||||
|
||||
# Hue shift: just verify it changes
|
||||
self._compare_param('grading_hue', {'grading_hue': 0.2})
|
||||
|
||||
# Gamma < 1: should brighten (raise values that are < 1)
|
||||
def check_gamma(base, result, _data):
|
||||
return float(result.mean()) > float(base.mean()), \
|
||||
f"baseline={float(base.mean()):.1f} gamma={float(result.mean()):.1f}"
|
||||
self._compare_param('grading_gamma', {'grading_gamma': 0.7}, check_gamma)
|
||||
|
||||
# Sharpness: just verify it changes
|
||||
self._compare_param('grading_sharpness', {'grading_sharpness': 0.8})
|
||||
|
||||
# Color temperature warm: red channel mean should increase relative to blue
|
||||
def check_warm(base, result, _data):
|
||||
base_r, _, base_b = self._channel_means(base)
|
||||
res_r, _, res_b = self._channel_means(result)
|
||||
base_rb = base_r - base_b
|
||||
res_rb = res_r - res_b
|
||||
return res_rb > base_rb, f"baseline R-B={base_rb:.1f} warm R-B={res_rb:.1f}"
|
||||
self._compare_param('grading_color_temp_warm', {'grading_color_temp': 3000}, check_warm)
|
||||
|
||||
# Color temperature cool: blue should increase relative to red
|
||||
def check_cool(base, result, _data):
|
||||
base_r, _, base_b = self._channel_means(base)
|
||||
res_r, _, res_b = self._channel_means(result)
|
||||
base_rb = base_r - base_b
|
||||
res_rb = res_r - res_b
|
||||
return res_rb < base_rb, f"baseline R-B={base_rb:.1f} cool R-B={res_rb:.1f}"
|
||||
self._compare_param('grading_color_temp_cool', {'grading_color_temp': 10000}, check_cool)
|
||||
|
||||
# Vignette: corners should be darker than baseline corners
|
||||
def check_vignette(base, result, _data):
|
||||
h, w = base.shape[:2]
|
||||
corner_size = h // 8
|
||||
base_corners = np.concatenate([
|
||||
base[:corner_size, :corner_size].flatten(),
|
||||
base[:corner_size, -corner_size:].flatten(),
|
||||
base[-corner_size:, :corner_size].flatten(),
|
||||
base[-corner_size:, -corner_size:].flatten(),
|
||||
])
|
||||
result_corners = np.concatenate([
|
||||
result[:corner_size, :corner_size].flatten(),
|
||||
result[:corner_size, -corner_size:].flatten(),
|
||||
result[-corner_size:, :corner_size].flatten(),
|
||||
result[-corner_size:, -corner_size:].flatten(),
|
||||
])
|
||||
return float(result_corners.mean()) < float(base_corners.mean()), \
|
||||
f"baseline_corners={float(base_corners.mean()):.1f} vignette_corners={float(result_corners.mean()):.1f}"
|
||||
self._compare_param('grading_vignette', {'grading_vignette': 0.8}, check_vignette)
|
||||
|
||||
# Grain: just verify it changes (stochastic)
|
||||
self._compare_param('grading_grain', {'grading_grain': 0.5})
|
||||
|
||||
# Shadows/midtones/highlights: verify changes
|
||||
self._compare_param('grading_shadows', {'grading_shadows': 0.5})
|
||||
self._compare_param('grading_highlights', {'grading_highlights': -0.3})
|
||||
|
||||
# CLAHE: should increase local contrast
|
||||
self._compare_param('grading_clahe', {'grading_clahe_clip': 2.0})
|
||||
|
||||
# Split toning: verify changes
|
||||
self._compare_param('grading_split_toning', {
|
||||
'grading_shadows_tint': '#003366',
|
||||
'grading_highlights_tint': '#ffcc00',
|
||||
})
|
||||
|
||||
# -- Correction params --
|
||||
|
||||
# Latent brightness: should change output and appear in metadata
|
||||
def check_correction_meta(key):
|
||||
def _check(_base, _result, data):
|
||||
info = self._get_info(data)
|
||||
return key in info, f"'{key}' {'found' if key in info else 'missing'} in info"
|
||||
return _check
|
||||
self._compare_param('hdr_brightness', {'hdr_brightness': 2.0}, check_correction_meta('Latent brightness'))
|
||||
self._compare_param('hdr_color', {'hdr_color': 0.8, 'hdr_mode': 1}, check_correction_meta('Latent color'))
|
||||
self._compare_param('hdr_sharpen', {'hdr_sharpen': 1.5}, check_correction_meta('Latent sharpen'))
|
||||
self._compare_param('hdr_clamp', {'hdr_clamp': True, 'hdr_threshold': 0.7}, check_correction_meta('Latent clamp'))
|
||||
|
||||
# Isolation: verify params from one request don't leak to the next
|
||||
data_after, _ = self._txt2img()
|
||||
arr_after = self._decode_image(data_after)
|
||||
baseline, _ = self._generate_baseline()
|
||||
if baseline is not None and arr_after is not None:
|
||||
leak_diff = self._pixel_diff(baseline, arr_after)
|
||||
no_leak = leak_diff < 0.5
|
||||
self.record(no_leak, 'param_isolation',
|
||||
f"post-grading baseline diff={leak_diff:.4f}" if no_leak
|
||||
else f"LEAK: baseline changed after grading requests (diff={leak_diff:.2f})")
|
||||
|
||||
# =========================================================================
|
||||
# Runner
|
||||
# =========================================================================
|
||||
|
||||
def run_all(self):
|
||||
print("=" * 60)
|
||||
print("Generation API Test Suite")
|
||||
print(f"Server: {self.base_url}")
|
||||
print(f"Steps: {self.steps}")
|
||||
print("=" * 60)
|
||||
|
||||
# Samplers
|
||||
available = self.test_samplers_list()
|
||||
self.test_samplers_generate(available)
|
||||
|
||||
# Grading
|
||||
self.run_grading_tests()
|
||||
|
||||
# Corrections
|
||||
self.run_correction_tests()
|
||||
|
||||
# Per-request param validation (baseline comparison)
|
||||
self.run_param_validation_tests()
|
||||
|
||||
# Summary
|
||||
print("\n" + "=" * 60)
|
||||
print("Results")
|
||||
print("=" * 60)
|
||||
total_passed = 0
|
||||
total_failed = 0
|
||||
total_skipped = 0
|
||||
for cat, data in self.results.items():
|
||||
total_passed += data['passed']
|
||||
total_failed += data['failed']
|
||||
total_skipped += data['skipped']
|
||||
status = 'PASS' if data['failed'] == 0 else 'FAIL'
|
||||
print(f" {cat}: {data['passed']} passed, {data['failed']} failed, {data['skipped']} skipped [{status}]")
|
||||
print(f" Total: {total_passed} passed, {total_failed} failed, {total_skipped} skipped")
|
||||
print("=" * 60)
|
||||
return total_failed == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description='Generation API Tests (samplers, grading, correction)')
|
||||
parser.add_argument('--url', default=os.environ.get('SDAPI_URL', 'http://127.0.0.1:7860'), help='server URL')
|
||||
parser.add_argument('--steps', type=int, default=10, help='generation steps (lower = faster tests)')
|
||||
args = parser.parse_args()
|
||||
test = GenerationAPITest(args.url, args.steps)
|
||||
success = test.run_all()
|
||||
sys.exit(0 if success else 1)
|
||||
@@ -0,0 +1,633 @@
|
||||
#!/usr/bin/env python
|
||||
"""
|
||||
Offline unit tests for color grading and latent corrections.
|
||||
|
||||
Tests two systems:
|
||||
- Pixel-space color grading (modules/processing_grading.py)
|
||||
- Latent-space corrections (modules/processing_correction.py)
|
||||
|
||||
No running server required. Tests core logic with synthetic inputs.
|
||||
|
||||
Usage:
|
||||
python test/test-grading.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import types
|
||||
import torch
|
||||
import numpy as np
|
||||
from types import SimpleNamespace
|
||||
|
||||
script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, script_dir)
|
||||
os.chdir(script_dir)
|
||||
|
||||
os.environ['SD_INSTALL_QUIET'] = '1'
|
||||
|
||||
# Initialize cmd_args before any module imports (required by shared.py)
|
||||
import modules.cmd_args
|
||||
import installer
|
||||
installer.add_args(modules.cmd_args.parser)
|
||||
modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([])
|
||||
|
||||
# Mock sd_vae_taesd to break circular import:
|
||||
# processing_correction -> sd_vae_taesd -> shared -> shared_items -> sd_vae_taesd (circle)
|
||||
_mock_taesd = types.ModuleType('modules.vae.sd_vae_taesd')
|
||||
_mock_taesd.TAESD_MODELS = {'taesd': None}
|
||||
_mock_taesd.CQYAN_MODELS = {}
|
||||
_mock_taesd.encode = lambda x: torch.zeros(1, 4, 8, 8)
|
||||
sys.modules['modules.vae.sd_vae_taesd'] = _mock_taesd
|
||||
|
||||
from modules.errors import log
|
||||
|
||||
# Results tracking
|
||||
results = {
|
||||
'grading_params': {'passed': 0, 'failed': 0, 'tests': []},
|
||||
'grading_functions': {'passed': 0, 'failed': 0, 'tests': []},
|
||||
'correction_primitives': {'passed': 0, 'failed': 0, 'tests': []},
|
||||
'correction_pipeline': {'passed': 0, 'failed': 0, 'tests': []},
|
||||
}
|
||||
current_category = 'grading_params'
|
||||
|
||||
|
||||
def record(passed, name, detail=''):
|
||||
status = 'PASS' if passed else 'FAIL'
|
||||
results[current_category]['passed' if passed else 'failed'] += 1
|
||||
results[current_category]['tests'].append((status, name))
|
||||
msg = f' {status}: {name}'
|
||||
if detail:
|
||||
msg += f' ({detail})'
|
||||
if passed:
|
||||
log.info(msg)
|
||||
else:
|
||||
log.error(msg)
|
||||
|
||||
|
||||
def set_category(cat):
|
||||
global current_category # pylint: disable=global-statement
|
||||
current_category = cat
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Color Grading: GradingParams and utility functions
|
||||
# ============================================================
|
||||
|
||||
def test_grading_params_defaults():
|
||||
"""GradingParams has correct defaults."""
|
||||
from modules.processing_grading import GradingParams
|
||||
p = GradingParams()
|
||||
assert p.brightness == 0.0
|
||||
assert p.contrast == 0.0
|
||||
assert p.saturation == 0.0
|
||||
assert p.hue == 0.0
|
||||
assert p.gamma == 1.0
|
||||
assert p.sharpness == 0.0
|
||||
assert p.color_temp == 6500
|
||||
assert p.shadows == 0.0
|
||||
assert p.midtones == 0.0
|
||||
assert p.highlights == 0.0
|
||||
assert p.clahe_clip == 0.0
|
||||
assert p.clahe_grid == 8
|
||||
assert p.shadows_tint == "#000000"
|
||||
assert p.highlights_tint == "#ffffff"
|
||||
assert p.split_tone_balance == 0.5
|
||||
assert p.vignette == 0.0
|
||||
assert p.grain == 0.0
|
||||
assert p.lut_file == ""
|
||||
assert p.lut_strength == 1.0
|
||||
return True
|
||||
|
||||
|
||||
def test_grading_is_active():
|
||||
"""is_active() returns False for defaults, True when any param differs."""
|
||||
from modules.processing_grading import GradingParams, is_active
|
||||
assert not is_active(GradingParams()), "defaults should be inactive"
|
||||
assert is_active(GradingParams(brightness=0.1)), "non-default brightness should be active"
|
||||
assert is_active(GradingParams(gamma=0.9)), "non-default gamma should be active"
|
||||
assert is_active(GradingParams(shadows_tint="#ff0000")), "non-default tint should be active"
|
||||
assert is_active(GradingParams(vignette=0.5)), "non-default vignette should be active"
|
||||
assert not is_active(GradingParams(brightness=0.0, gamma=1.0, color_temp=6500)), "all-default should be inactive"
|
||||
return True
|
||||
|
||||
|
||||
def test_grading_float_coercion():
|
||||
"""__post_init__ coerces int inputs to float (Gradio sends int for float sliders)."""
|
||||
from modules.processing_grading import GradingParams
|
||||
p = GradingParams(brightness=1, contrast=2, gamma=1, color_temp=6500)
|
||||
assert isinstance(p.brightness, float), f"expected float, got {type(p.brightness)}"
|
||||
assert isinstance(p.contrast, float), f"expected float, got {type(p.contrast)}"
|
||||
assert isinstance(p.gamma, float), f"expected float, got {type(p.gamma)}"
|
||||
assert isinstance(p.color_temp, float), f"expected float, got {type(p.color_temp)}"
|
||||
return True
|
||||
|
||||
|
||||
def test_hex_to_rgb():
|
||||
"""_hex_to_rgb converts hex color strings correctly."""
|
||||
from modules.processing_grading import _hex_to_rgb
|
||||
assert _hex_to_rgb("#000000") == (0.0, 0.0, 0.0), "black"
|
||||
assert _hex_to_rgb("#ffffff") == (1.0, 1.0, 1.0), "white"
|
||||
r, g, b = _hex_to_rgb("#ff0000")
|
||||
assert abs(r - 1.0) < 1e-6 and abs(g) < 1e-6 and abs(b) < 1e-6, "red"
|
||||
r, g, b = _hex_to_rgb("#00ff00")
|
||||
assert abs(r) < 1e-6 and abs(g - 1.0) < 1e-6 and abs(b) < 1e-6, "green"
|
||||
r, g, b = _hex_to_rgb("#0000ff")
|
||||
assert abs(r) < 1e-6 and abs(g) < 1e-6 and abs(b - 1.0) < 1e-6, "blue"
|
||||
# without hash
|
||||
r, g, b = _hex_to_rgb("ff8040")
|
||||
assert r > g > b, "orange-ish ordering"
|
||||
# invalid length returns black
|
||||
assert _hex_to_rgb("#fff") == (0.0, 0.0, 0.0), "short hex returns black"
|
||||
return True
|
||||
|
||||
|
||||
def test_kelvin_to_rgb():
|
||||
"""_kelvin_to_rgb_scale returns sensible values at known temperatures."""
|
||||
from modules.processing_grading import _kelvin_to_rgb_scale
|
||||
# 6500K (reference) should be approximately (1, 1, 1) - tolerance is wide because
|
||||
# the Planckian formula approximation normalizes to a hardcoded ref point
|
||||
r, g, b = _kelvin_to_rgb_scale(6500)
|
||||
assert abs(r - 1.0) < 0.15 and abs(g - 1.0) < 0.15 and abs(b - 1.0) < 0.15, f"6500K: ({r:.3f}, {g:.3f}, {b:.3f})"
|
||||
# warm (3000K) should have r > b
|
||||
r, g, b = _kelvin_to_rgb_scale(3000)
|
||||
assert r > b, f"3000K should be warm: r={r:.3f} b={b:.3f}"
|
||||
# cool (10000K) should have b > r
|
||||
r, g, b = _kelvin_to_rgb_scale(10000)
|
||||
assert b > r, f"10000K should be cool: r={r:.3f} b={b:.3f}"
|
||||
# all values should be non-negative; very low temps have zero blue (physically correct)
|
||||
for temp in [1000, 2000, 4000, 8000, 15000, 40000]:
|
||||
r, g, b = _kelvin_to_rgb_scale(temp)
|
||||
assert r >= 0 and g >= 0 and b >= 0, f"{temp}K has negative channel: ({r:.3f}, {g:.3f}, {b:.3f})"
|
||||
# moderate temps should have all positive channels
|
||||
for temp in [3000, 5000, 6500, 10000]:
|
||||
r, g, b = _kelvin_to_rgb_scale(temp)
|
||||
assert r > 0 and g > 0 and b > 0, f"{temp}K has non-positive channel: ({r:.3f}, {g:.3f}, {b:.3f})"
|
||||
return True
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Color Grading: torch-based functions (need kornia for some)
|
||||
# ============================================================
|
||||
|
||||
def _make_test_image_tensor(h=64, w=64):
|
||||
"""Create a synthetic RGB test image tensor [1, 3, H, W] in [0, 1]."""
|
||||
torch.manual_seed(42)
|
||||
return torch.rand(1, 3, h, w, dtype=torch.float32)
|
||||
|
||||
|
||||
def _make_test_pil_image(h=64, w=64):
|
||||
"""Create a synthetic RGB PIL image."""
|
||||
from PIL import Image
|
||||
arr = np.random.RandomState(42).randint(0, 255, (h, w, 3), dtype=np.uint8)
|
||||
return Image.fromarray(arr, 'RGB')
|
||||
|
||||
|
||||
def test_apply_vignette():
|
||||
"""_apply_vignette darkens edges more than center."""
|
||||
from modules.processing_grading import _apply_vignette
|
||||
img = torch.ones(1, 3, 64, 64, dtype=torch.float32)
|
||||
result = _apply_vignette(img, strength=1.0)
|
||||
center_val = result[0, 0, 32, 32].item()
|
||||
corner_val = result[0, 0, 0, 0].item()
|
||||
assert center_val > corner_val, f"center ({center_val:.3f}) should be brighter than corner ({corner_val:.3f})"
|
||||
assert result.shape == img.shape, "shape preserved"
|
||||
assert not torch.isnan(result).any(), "no NaN"
|
||||
# zero strength should be identity
|
||||
result_zero = _apply_vignette(img, strength=0.0)
|
||||
assert torch.allclose(result_zero, img), "zero strength is identity"
|
||||
return True
|
||||
|
||||
|
||||
def test_apply_grain():
|
||||
"""_apply_grain adds noise (output differs from input, stays in valid range)."""
|
||||
from modules.processing_grading import _apply_grain
|
||||
img = torch.ones(1, 3, 64, 64, dtype=torch.float32) * 0.5
|
||||
result = _apply_grain(img, strength=0.5)
|
||||
assert not torch.equal(result, img), "grain should modify the image"
|
||||
assert result.shape == img.shape, "shape preserved"
|
||||
assert result.min() >= 0.0 and result.max() <= 1.0, "output clamped to [0, 1]"
|
||||
assert not torch.isnan(result).any(), "no NaN"
|
||||
return True
|
||||
|
||||
|
||||
def test_apply_color_temp():
|
||||
"""_apply_color_temp shifts R/B channels for warm/cool temperatures."""
|
||||
from modules.processing_grading import _apply_color_temp
|
||||
img = torch.ones(1, 3, 64, 64, dtype=torch.float32) * 0.5
|
||||
# warm
|
||||
warm = _apply_color_temp(img, 3000)
|
||||
assert warm[0, 0].mean() > warm[0, 2].mean(), "warm: red > blue"
|
||||
# cool
|
||||
cool = _apply_color_temp(img, 10000)
|
||||
assert cool[0, 2].mean() > cool[0, 0].mean(), "cool: blue > red"
|
||||
# neutral
|
||||
neutral = _apply_color_temp(img, 6500)
|
||||
assert torch.allclose(neutral, img, atol=0.05), "6500K is near-neutral"
|
||||
assert warm.shape == img.shape, "shape preserved"
|
||||
return True
|
||||
|
||||
|
||||
def test_apply_shadows_midtones_highlights():
|
||||
"""_apply_shadows_midtones_highlights modifies tone without NaN/shape issues."""
|
||||
try:
|
||||
from modules.processing_grading import _apply_shadows_midtones_highlights
|
||||
except ImportError:
|
||||
return None # kornia not available
|
||||
img = _make_test_image_tensor()
|
||||
# shadows boost
|
||||
result = _apply_shadows_midtones_highlights(img, shadows=0.5, midtones=0.0, highlights=0.0)
|
||||
assert result.shape == img.shape, "shape preserved"
|
||||
assert not torch.isnan(result).any(), "no NaN"
|
||||
assert result.min() >= 0.0 and result.max() <= 1.0, "output in [0, 1]"
|
||||
# all zero should be near-identity (kornia conversions may introduce tiny diffs)
|
||||
result_zero = _apply_shadows_midtones_highlights(img, shadows=0.0, midtones=0.0, highlights=0.0)
|
||||
assert torch.allclose(result_zero, img, atol=1e-3), "zero params is near-identity"
|
||||
return True
|
||||
|
||||
|
||||
def test_grade_image_pipeline():
|
||||
"""Full grade_image pipeline runs without errors for various param combos."""
|
||||
try:
|
||||
import modules.devices as devices_mod
|
||||
devices_mod.device = torch.device('cpu')
|
||||
devices_mod.dtype = torch.float32
|
||||
from modules.processing_grading import GradingParams, grade_image, is_active
|
||||
except ImportError:
|
||||
return None # kornia not available
|
||||
img = _make_test_pil_image()
|
||||
# basic adjustments
|
||||
params = GradingParams(brightness=0.1, contrast=0.2, saturation=-0.1)
|
||||
assert is_active(params)
|
||||
result = grade_image(img, params)
|
||||
assert result.size == img.size, "output size matches input"
|
||||
assert result.mode == 'RGB', "output is RGB"
|
||||
# tone adjustments
|
||||
params = GradingParams(shadows=0.3, midtones=-0.2, highlights=0.1)
|
||||
result = grade_image(img, params)
|
||||
assert result.size == img.size
|
||||
# effects
|
||||
params = GradingParams(vignette=0.5, grain=0.3)
|
||||
result = grade_image(img, params)
|
||||
assert result.size == img.size
|
||||
# hue and gamma
|
||||
params = GradingParams(hue=0.1, gamma=0.8, sharpness=0.5)
|
||||
result = grade_image(img, params)
|
||||
assert result.size == img.size
|
||||
# color temp
|
||||
params = GradingParams(color_temp=3000)
|
||||
result = grade_image(img, params)
|
||||
assert result.size == img.size
|
||||
# split toning
|
||||
params = GradingParams(shadows_tint="#003366", highlights_tint="#ffcc00", split_tone_balance=0.7)
|
||||
result = grade_image(img, params)
|
||||
assert result.size == img.size
|
||||
return True
|
||||
|
||||
|
||||
def test_grade_image_edge_cases():
|
||||
"""grade_image handles edge cases: all-black, all-white, tiny images."""
|
||||
try:
|
||||
import modules.devices as devices_mod
|
||||
devices_mod.device = torch.device('cpu')
|
||||
devices_mod.dtype = torch.float32
|
||||
from modules.processing_grading import GradingParams, grade_image
|
||||
except ImportError:
|
||||
return None
|
||||
from PIL import Image
|
||||
params = GradingParams(brightness=0.2, contrast=0.3, vignette=0.5, grain=0.2)
|
||||
# all black
|
||||
black = Image.fromarray(np.zeros((64, 64, 3), dtype=np.uint8), 'RGB')
|
||||
result = grade_image(black, params)
|
||||
assert result.size == black.size, "black image handled"
|
||||
# all white
|
||||
white = Image.fromarray(np.full((64, 64, 3), 255, dtype=np.uint8), 'RGB')
|
||||
result = grade_image(white, params)
|
||||
assert result.size == white.size, "white image handled"
|
||||
# tiny image
|
||||
tiny = Image.fromarray(np.random.randint(0, 255, (4, 4, 3), dtype=np.uint8), 'RGB')
|
||||
result = grade_image(tiny, params)
|
||||
assert result.size == tiny.size, "tiny image handled"
|
||||
return True
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Latent Corrections: primitive tensor operations
|
||||
# ============================================================
|
||||
|
||||
def test_soft_clamp_tensor():
|
||||
"""soft_clamp_tensor shrinks outliers toward mean, preserves values within bounds."""
|
||||
from modules.processing_correction import soft_clamp_tensor
|
||||
# within bounds: no change
|
||||
tensor = torch.randn(4, 64, 64) * 0.5
|
||||
result = soft_clamp_tensor(tensor, threshold=0.8, boundary=4)
|
||||
assert torch.allclose(result, tensor, atol=1e-5), "within-bounds tensor unchanged"
|
||||
# with outliers: should clamp
|
||||
tensor_outliers = torch.randn(4, 64, 64)
|
||||
tensor_outliers[0, 0, 0] = 10.0
|
||||
tensor_outliers[1, 0, 0] = -10.0
|
||||
result = soft_clamp_tensor(tensor_outliers.clone(), threshold=0.8, boundary=4)
|
||||
assert result[0, 0, 0] < tensor_outliers[0, 0, 0], "positive outlier reduced"
|
||||
assert result[1, 0, 0] > tensor_outliers[1, 0, 0], "negative outlier raised"
|
||||
assert result.shape == tensor_outliers.shape, "shape preserved"
|
||||
assert not torch.isnan(result).any(), "no NaN"
|
||||
# zero threshold: identity
|
||||
result_zero = soft_clamp_tensor(tensor_outliers.clone(), threshold=0, boundary=4)
|
||||
assert torch.allclose(result_zero, tensor_outliers), "zero threshold is identity"
|
||||
return True
|
||||
|
||||
|
||||
def test_center_tensor():
|
||||
"""center_tensor adjusts mean of tensor channels."""
|
||||
from modules.processing_correction import center_tensor
|
||||
tensor = torch.randn(4, 64, 64) + 2.0 # offset mean
|
||||
original_mean = tensor.mean().item()
|
||||
# full shift should reduce mean toward offset
|
||||
result = center_tensor(tensor.clone(), channel_shift=0.0, full_shift=1.0, offset=0.0)
|
||||
assert abs(result.mean().item()) < abs(original_mean), "full shift centers toward zero"
|
||||
# channel shift
|
||||
result_ch = center_tensor(tensor.clone(), channel_shift=1.0, full_shift=0.0, offset=0.0)
|
||||
for c in range(4):
|
||||
assert abs(result_ch[c].mean().item()) < abs(tensor[c].mean().item()), f"channel {c} centered"
|
||||
# no-op
|
||||
result_noop = center_tensor(tensor.clone(), channel_shift=0.0, full_shift=0.0, offset=0.0)
|
||||
assert torch.allclose(result_noop, tensor), "zero params is identity"
|
||||
# with offset
|
||||
result_offset = center_tensor(tensor.clone(), channel_shift=0.0, full_shift=1.0, offset=5.0)
|
||||
assert result_offset.mean().item() > 0, "offset shifts mean positive"
|
||||
return True
|
||||
|
||||
|
||||
def test_sharpen_tensor():
|
||||
"""sharpen_tensor applies sharpening convolution, preserves shape."""
|
||||
from modules.processing_correction import sharpen_tensor
|
||||
tensor = torch.randn(4, 64, 64)
|
||||
# zero ratio: identity
|
||||
result_zero = sharpen_tensor(tensor.clone(), ratio=0)
|
||||
assert torch.allclose(result_zero, tensor), "zero ratio is identity"
|
||||
# positive ratio: should modify
|
||||
result = sharpen_tensor(tensor.clone(), ratio=0.5)
|
||||
assert result.shape == tensor.shape, "shape preserved"
|
||||
assert not torch.isnan(result).any(), "no NaN"
|
||||
assert not torch.isinf(result).any(), "no Inf"
|
||||
assert not torch.equal(result, tensor), "sharpening modifies tensor"
|
||||
return True
|
||||
|
||||
|
||||
def test_maximize_tensor():
|
||||
"""maximize_tensor normalizes tensor range."""
|
||||
from modules.processing_correction import maximize_tensor
|
||||
tensor = torch.randn(4, 64, 64) * 0.5
|
||||
# boundary 1.0: identity
|
||||
result_id = maximize_tensor(tensor.clone(), boundary=1.0)
|
||||
assert torch.allclose(result_id, tensor), "boundary 1.0 is identity"
|
||||
# boundary 2.0: should expand range
|
||||
result = maximize_tensor(tensor.clone(), boundary=2.0)
|
||||
assert result.abs().max() > tensor.abs().max(), "boundary 2.0 expands range"
|
||||
assert result.shape == tensor.shape, "shape preserved"
|
||||
assert not torch.isnan(result).any(), "no NaN"
|
||||
# boundary 0.5: should compress range
|
||||
result_small = maximize_tensor(tensor.clone(), boundary=0.5)
|
||||
assert result_small.abs().max() < tensor.abs().max() + 0.1, "boundary 0.5 compresses"
|
||||
return True
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Latent Corrections: correction() pipeline with mock p object
|
||||
# ============================================================
|
||||
|
||||
def _make_mock_p(**overrides):
|
||||
"""Create a mock processing object with default hdr params."""
|
||||
defaults = {
|
||||
'hdr_mode': 0,
|
||||
'hdr_brightness': 0.0,
|
||||
'hdr_color': 0.0,
|
||||
'hdr_sharpen': 0.0,
|
||||
'hdr_clamp': False,
|
||||
'hdr_boundary': 4.0,
|
||||
'hdr_threshold': 0.95,
|
||||
'hdr_maximize': False,
|
||||
'hdr_max_center': 0.6,
|
||||
'hdr_max_boundary': 1.0,
|
||||
'hdr_color_picker': '#000000',
|
||||
'hdr_tint_ratio': 0.0,
|
||||
'correction_total_steps': 20,
|
||||
'correction_steps_mid': 10,
|
||||
'correction_steps_late': 4,
|
||||
'extra_generation_params': {},
|
||||
}
|
||||
defaults.update(overrides)
|
||||
return SimpleNamespace(**defaults)
|
||||
|
||||
|
||||
def test_correction_noop():
|
||||
"""correction() with all-zero params is near-identity."""
|
||||
from modules.processing_correction import correction
|
||||
p = _make_mock_p()
|
||||
latent = torch.randn(4, 64, 64)
|
||||
for step in [0, 5, 10, 15, 19]:
|
||||
result = correction(p, 500, latent.clone(), step=step)
|
||||
assert torch.allclose(result, latent, atol=1e-5), f"step {step}: no-op correction should be identity"
|
||||
return True
|
||||
|
||||
|
||||
def test_correction_early_clamp():
|
||||
"""correction() applies soft_clamp in early steps when hdr_clamp=True."""
|
||||
from modules.processing_correction import correction
|
||||
p = _make_mock_p(hdr_clamp=True, hdr_threshold=0.8, hdr_boundary=4.0)
|
||||
latent = torch.randn(4, 64, 64)
|
||||
latent[0, 0, 0] = 10.0 # outlier
|
||||
# step 0 of 20 = progress 0.0 (early)
|
||||
result = correction(p, 999, latent.clone(), step=0)
|
||||
assert result[0, 0, 0] < 10.0, "outlier should be clamped"
|
||||
assert "Latent clamp" in p.extra_generation_params, "clamp recorded in params"
|
||||
return True
|
||||
|
||||
|
||||
def test_correction_mid_color():
|
||||
"""correction() applies color centering in mid steps."""
|
||||
from modules.processing_correction import correction
|
||||
p = _make_mock_p(hdr_color=0.5)
|
||||
latent = torch.randn(4, 64, 64) + 1.0 # offset channels
|
||||
original_ch_means = [latent[c].mean().item() for c in range(1, 4)]
|
||||
# step 6 of 20 = progress 0.3 (mid range)
|
||||
result = correction(p, 700, latent.clone(), step=6)
|
||||
new_ch_means = [result[c].mean().item() for c in range(1, 4)]
|
||||
# at least some channels should have their mean reduced (centered)
|
||||
centered_count = sum(1 for o, n in zip(original_ch_means, new_ch_means) if abs(n) < abs(o))
|
||||
assert centered_count > 0, "at least one color channel should be more centered"
|
||||
assert "Latent color" in p.extra_generation_params, "color recorded in params"
|
||||
return True
|
||||
|
||||
|
||||
def test_correction_late_brightness():
|
||||
"""correction() applies brightness offset in late steps."""
|
||||
from modules.processing_correction import correction
|
||||
p = _make_mock_p(hdr_brightness=2.0)
|
||||
latent = torch.randn(4, 64, 64)
|
||||
original_mean = latent[0].mean().item()
|
||||
# step 17 of 20 = progress 0.85 (late)
|
||||
result = correction(p, 100, latent.clone(), step=17)
|
||||
new_mean = result[0].mean().item()
|
||||
assert new_mean != original_mean, "brightness should shift channel 0 mean"
|
||||
assert "Latent brightness" in p.extra_generation_params, "brightness recorded in params"
|
||||
return True
|
||||
|
||||
|
||||
def test_correction_sharpen():
|
||||
"""correction() applies sharpening in sharpen range."""
|
||||
from modules.processing_correction import correction
|
||||
p = _make_mock_p(hdr_sharpen=1.0)
|
||||
latent = torch.randn(4, 64, 64)
|
||||
# step 15 of 20 = progress 0.75 (sharpen range)
|
||||
result = correction(p, 200, latent.clone(), step=15)
|
||||
assert not torch.equal(result, latent), "sharpening should modify latent"
|
||||
assert "Latent sharpen" in p.extra_generation_params, "sharpen recorded in params"
|
||||
return True
|
||||
|
||||
|
||||
def test_correction_maximize():
|
||||
"""correction() applies maximize/normalize in very late steps."""
|
||||
from modules.processing_correction import correction
|
||||
p = _make_mock_p(hdr_maximize=True, hdr_max_center=0.6, hdr_max_boundary=2.0)
|
||||
latent = torch.randn(4, 64, 64) * 0.5
|
||||
# step 19 of 20 = progress 0.95 (very late)
|
||||
result = correction(p, 10, latent.clone(), step=19)
|
||||
assert result.abs().max() > latent.abs().max(), "maximize should expand range"
|
||||
assert "Latent max" in p.extra_generation_params, "maximize recorded in params"
|
||||
return True
|
||||
|
||||
|
||||
def test_correction_multichannel():
|
||||
"""correction() uses multi-channel path for >4 channel latents."""
|
||||
from modules.processing_correction import correction
|
||||
p = _make_mock_p(hdr_brightness=2.0, hdr_color=0.5)
|
||||
# 16-channel latent (e.g. Flux 2)
|
||||
latent = torch.randn(16, 64, 64)
|
||||
# mid step: color centering on all channels
|
||||
p.extra_generation_params = {}
|
||||
result_mid = correction(p, 700, latent.clone(), step=6)
|
||||
assert result_mid.shape == latent.shape, "multi-channel shape preserved"
|
||||
assert not torch.isnan(result_mid).any(), "no NaN"
|
||||
# late step: brightness via multiplicative scaling
|
||||
p.extra_generation_params = {}
|
||||
result_late = correction(p, 100, latent.clone(), step=17)
|
||||
assert result_late.shape == latent.shape, "multi-channel shape preserved"
|
||||
assert not torch.isnan(result_late).any(), "no NaN"
|
||||
assert "Latent brightness" in p.extra_generation_params, "brightness recorded"
|
||||
return True
|
||||
|
||||
|
||||
def test_correction_shape_preservation():
|
||||
"""correction() preserves shape and dtype for various latent sizes."""
|
||||
from modules.processing_correction import correction
|
||||
p = _make_mock_p(hdr_clamp=True, hdr_color=0.3, hdr_brightness=1.0, hdr_sharpen=0.5)
|
||||
shapes = [(4, 64, 64), (4, 32, 32), (4, 128, 128), (8, 64, 64), (16, 32, 32)]
|
||||
for shape in shapes:
|
||||
for step in [0, 6, 15, 19]:
|
||||
p.extra_generation_params = {}
|
||||
latent = torch.randn(shape)
|
||||
result = correction(p, 500, latent.clone(), step=step)
|
||||
assert result.shape == latent.shape, f"shape {shape} step {step}: shape mismatch"
|
||||
assert result.dtype == latent.dtype, f"shape {shape} step {step}: dtype mismatch"
|
||||
assert not torch.isnan(result).any(), f"shape {shape} step {step}: NaN"
|
||||
assert not torch.isinf(result).any(), f"shape {shape} step {step}: Inf"
|
||||
return True
|
||||
|
||||
|
||||
def test_correction_step_ranges():
|
||||
"""correction() applies different operations at different progress points."""
|
||||
from modules.processing_correction import correction
|
||||
p = _make_mock_p(
|
||||
hdr_clamp=True, hdr_color=0.5, hdr_brightness=1.0,
|
||||
hdr_sharpen=0.5, hdr_maximize=True, hdr_max_boundary=2.0,
|
||||
)
|
||||
latent = torch.randn(4, 64, 64)
|
||||
latent[0, 0, 0] = 10.0 # outlier for clamp testing
|
||||
expected_params_per_range = {
|
||||
0: ["Latent clamp"], # early: progress 0.0
|
||||
6: ["Latent color"], # mid: progress 0.3
|
||||
15: ["Latent sharpen"], # sharpen: progress 0.75
|
||||
17: ["Latent brightness"], # late: progress 0.85
|
||||
19: ["Latent max"], # very late: progress 0.95
|
||||
}
|
||||
for step, expected_keys in expected_params_per_range.items():
|
||||
p.extra_generation_params = {}
|
||||
correction(p, 500, latent.clone(), step=step)
|
||||
for key in expected_keys:
|
||||
assert key in p.extra_generation_params, f"step {step}: expected '{key}' in params, got {list(p.extra_generation_params.keys())}"
|
||||
return True
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Test runner
|
||||
# ============================================================
|
||||
|
||||
def run_test(fn):
|
||||
name = fn.__name__
|
||||
try:
|
||||
result = fn()
|
||||
if result is None:
|
||||
log.warning(f' SKIP: {name} (dependency not available)')
|
||||
return
|
||||
record(True, name)
|
||||
except AssertionError as e:
|
||||
record(False, name, str(e))
|
||||
except Exception as e:
|
||||
record(False, name, f"exception: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
def run_tests():
|
||||
t0 = time.time()
|
||||
|
||||
# Grading params (pure Python, no GPU deps)
|
||||
set_category('grading_params')
|
||||
log.warning('=== Color Grading: Params & Utilities ===')
|
||||
for fn in [test_grading_params_defaults, test_grading_is_active, test_grading_float_coercion,
|
||||
test_hex_to_rgb, test_kelvin_to_rgb]:
|
||||
run_test(fn)
|
||||
|
||||
# Grading functions (need torch, some need kornia)
|
||||
set_category('grading_functions')
|
||||
log.warning('=== Color Grading: Tensor Operations ===')
|
||||
for fn in [test_apply_vignette, test_apply_grain, test_apply_color_temp,
|
||||
test_apply_shadows_midtones_highlights, test_grade_image_pipeline,
|
||||
test_grade_image_edge_cases]:
|
||||
run_test(fn)
|
||||
|
||||
# Correction primitives (pure torch)
|
||||
set_category('correction_primitives')
|
||||
log.warning('=== Latent Corrections: Primitives ===')
|
||||
for fn in [test_soft_clamp_tensor, test_center_tensor, test_sharpen_tensor,
|
||||
test_maximize_tensor]:
|
||||
run_test(fn)
|
||||
|
||||
# Correction pipeline (mock p object)
|
||||
set_category('correction_pipeline')
|
||||
log.warning('=== Latent Corrections: Pipeline ===')
|
||||
for fn in [test_correction_noop, test_correction_early_clamp, test_correction_mid_color,
|
||||
test_correction_late_brightness, test_correction_sharpen, test_correction_maximize,
|
||||
test_correction_multichannel, test_correction_shape_preservation,
|
||||
test_correction_step_ranges]:
|
||||
run_test(fn)
|
||||
|
||||
t1 = time.time()
|
||||
|
||||
# Summary
|
||||
log.warning('=== Results ===')
|
||||
total_passed = 0
|
||||
total_failed = 0
|
||||
for cat, data in results.items():
|
||||
total_passed += data['passed']
|
||||
total_failed += data['failed']
|
||||
status = 'PASS' if data['failed'] == 0 else 'FAIL'
|
||||
log.info(f' {cat}: {data["passed"]} passed, {data["failed"]} failed [{status}]')
|
||||
log.warning(f'Total: {total_passed} passed, {total_failed} failed in {t1 - t0:.2f}s')
|
||||
if total_failed > 0:
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_tests()
|
||||
Reference in New Issue
Block a user