Files
sglang/python/sglang/multimodal_gen/test/server/test_server_utils.py
blahblah 0a7a2017a0 [diffusion] refactor: refactor and simplify teacache for cachabledit and wanvideo (#16396)
Co-authored-by: Brain97 <Brain97@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
Co-authored-by: blahblah <blahblah>
2026-01-21 15:42:45 +08:00

1168 lines
39 KiB
Python

"""
Server management and performance validation for diffusion tests.
"""
from __future__ import annotations
import base64
import os
import shlex
import subprocess
import sys
import tempfile
import threading
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Sequence
from urllib.request import urlopen
import pytest
from openai import Client, OpenAI
from sglang.multimodal_gen.benchmarks.compare_perf import calculate_upper_bound
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.common import kill_process_tree
from sglang.multimodal_gen.runtime.utils.logging_utils import (
globally_suppress_loggers,
init_logger,
)
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord
from sglang.multimodal_gen.test.server.testcase_configs import (
DiffusionSamplingParams,
PerformanceSummary,
ScenarioConfig,
ToleranceConfig,
)
from sglang.multimodal_gen.test.slack_utils import upload_file_to_slack
from sglang.multimodal_gen.test.test_utils import (
get_expected_image_format,
is_image_url,
prepare_perf_log,
validate_image,
validate_image_file,
validate_openai_video,
validate_video_file,
)
logger = init_logger(__name__)
globally_suppress_loggers()
def download_image_from_url(url: str) -> Path:
"""Download an image from a URL to a temporary file.
Args:
url: The URL of the image to download
Returns:
Path to the downloaded temporary file
"""
logger.info(f"Downloading image from URL: {url}")
# Determine file extension from URL
ext = ".jpg" # default
if url.lower().endswith((".png", ".jpeg", ".jpg", ".webp", ".gif")):
ext = url[url.rfind(".") :]
# Create temporary file
temp_file = (
Path(tempfile.gettempdir()) / f"diffusion_test_image_{int(time.time())}{ext}"
)
try:
with urlopen(url, timeout=30) as response:
temp_file.write_bytes(response.read())
logger.info(f"Downloaded image to: {temp_file}")
return temp_file
except Exception as e:
logger.error(f"Failed to download image from {url}: {e}")
raise
def parse_dimensions(size_string: str | None) -> tuple[int | None, int | None]:
"""Parse a size string in "widthxheight" format to (width, height) tuple.
Args:
size_string: Size string in "widthxheight" format (e.g., "1024x1024") or None.
Spaces are automatically stripped.
Returns:
Tuple of (width, height) as integers if parsing succeeds, (None, None) otherwise.
"""
if not size_string:
return (None, None)
# Strip spaces from the entire string
size_string = size_string.strip()
if not size_string:
return (None, None)
# Split by "x"
parts = size_string.split("x")
if len(parts) != 2:
return (None, None)
# Strip spaces from each part and try to convert to int
try:
width_str = parts[0].strip()
height_str = parts[1].strip()
if not width_str or not height_str:
return (None, None)
width = int(width_str)
height = int(height_str)
# Validate that both are positive
if width <= 0 or height <= 0:
return (None, None)
return (width, height)
except ValueError:
return (None, None)
@dataclass
class ServerContext:
"""Context for a running diffusion server."""
port: int
process: subprocess.Popen
model: str
stdout_file: Path
perf_log_path: Path
log_dir: Path
_stdout_fh: Any = field(repr=False)
_log_thread: threading.Thread | None = field(default=None, repr=False)
def cleanup(self) -> None:
"""Clean up server resources."""
try:
kill_process_tree(self.process.pid)
except Exception:
pass
try:
self._stdout_fh.flush()
self._stdout_fh.close()
except Exception:
pass
# ROCm/AMD: Extra cleanup to ensure GPU memory is released between tests
# This is needed because ROCm memory release can be slower than CUDA
if current_platform.is_hip():
self._cleanup_rocm_gpu_memory()
# Clean up downloaded models if HF cache is not persistent
# This prevents disk exhaustion in CI when cache is not mounted
self._cleanup_hf_cache_if_not_persistent()
def _cleanup_hf_cache_if_not_persistent(self) -> None:
"""Clean up HF cache if it's not on a persistent volume.
When running in CI without persistent cache, downloaded models accumulate
and can cause disk/memory exhaustion. This cleans up the model after each
test if the cache is not persistent.
"""
import shutil
hf_home = os.environ.get("HF_HOME", "")
if not hf_home:
return
hf_hub_cache = os.path.join(hf_home, "hub")
# Check if HF cache is on a persistent volume by looking for a marker file
# or checking if the directory existed before this test run
persistent_marker = os.path.join(hf_home, ".persistent_cache")
if os.path.exists(persistent_marker):
logger.info("HF cache is persistent, skipping cleanup")
return
# Check if the cache directory is empty or was just created
# If it has very few models, it's likely not persistent
if not os.path.exists(hf_hub_cache):
return
try:
# Get model cache directories
model_dirs = [
d
for d in os.listdir(hf_hub_cache)
if d.startswith("models--")
and os.path.isdir(os.path.join(hf_hub_cache, d))
]
# If there are cached models but no persistent marker, clean up
# to prevent disk exhaustion in CI
if model_dirs:
logger.info(
"HF cache appears non-persistent (no .persistent_cache marker), "
"cleaning up %d model(s) to prevent disk exhaustion",
len(model_dirs),
)
for model_dir in model_dirs:
model_path = os.path.join(hf_hub_cache, model_dir)
try:
shutil.rmtree(model_path)
logger.info("Cleaned up model cache: %s", model_dir)
except Exception as e:
logger.warning("Failed to clean up %s: %s", model_dir, e)
except Exception as e:
logger.warning("Error during HF cache cleanup: %s", e)
def _cleanup_rocm_gpu_memory(self) -> None:
"""ROCm-specific cleanup to ensure GPU memory is fully released."""
import gc
# Wait for process to fully terminate
try:
self.process.wait(timeout=30)
except Exception:
pass
# Force garbage collection multiple times
for _ in range(3):
gc.collect()
# Clear HIP memory on all GPUs
try:
import torch
for i in range(torch.cuda.device_count()):
with torch.cuda.device(i):
torch.cuda.empty_cache()
torch.cuda.synchronize()
except Exception:
pass
# Wait for GPU memory to be released (ROCm can be much slower than CUDA)
# The GPU driver needs time to reclaim memory from killed processes
time.sleep(15)
class ServerManager:
"""Manages diffusion server lifecycle."""
def __init__(
self,
model: str,
port: int,
wait_deadline: float = 1200.0,
extra_args: str = "",
env_vars: dict[str, str] | None = None,
):
self.model = model
self.port = port
self.wait_deadline = wait_deadline
self.extra_args = extra_args
self.env_vars = env_vars or {}
def _wait_for_rocm_gpu_memory_clear(self, max_wait: float = 60.0) -> None:
"""ROCm-specific: Wait for GPU memory to be mostly free before starting.
ROCm GPU memory release from killed processes can be significantly slower
than CUDA, so we need to wait longer and be more patient.
"""
try:
import torch
if not torch.cuda.is_available():
return
start_time = time.time()
last_total_used = float("inf")
while time.time() - start_time < max_wait:
# Check GPU memory usage
total_used = 0
for i in range(torch.cuda.device_count()):
mem_info = torch.cuda.mem_get_info(i)
free, total = mem_info
used = total - free
total_used += used
# If less than 5GB is used across all GPUs, we're good
if total_used < 5 * 1024 * 1024 * 1024: # 5GB
logger.info(
"[server-test] ROCm GPU memory is clear (used: %.2f GB)",
total_used / (1024**3),
)
return
# Log progress
elapsed = int(time.time() - start_time)
if total_used < last_total_used:
logger.info(
"[server-test] ROCm: GPU memory clearing (used: %.2f GB, elapsed: %ds)",
total_used / (1024**3),
elapsed,
)
else:
logger.info(
"[server-test] ROCm: Waiting for GPU memory (used: %.2f GB, elapsed: %ds)",
total_used / (1024**3),
elapsed,
)
last_total_used = total_used
time.sleep(3)
# Final warning with detailed GPU info
logger.warning(
"[server-test] ROCm GPU memory not fully cleared after %.0fs (used: %.2f GB). "
"Proceeding anyway - this may cause OOM.",
max_wait,
total_used / (1024**3),
)
except Exception as e:
logger.debug("[server-test] Could not check ROCm GPU memory: %s", e)
def start(self) -> ServerContext:
"""Start the diffusion server and wait for readiness."""
# ROCm/AMD: Wait for GPU memory to be clear before starting
# This prevents OOM when running sequential tests on ROCm
if current_platform.is_hip():
self._wait_for_rocm_gpu_memory_clear()
log_dir, perf_log_path = prepare_perf_log()
safe_model_name = self.model.replace("/", "_")
stdout_path = (
Path(tempfile.gettempdir())
/ f"sgl_server_{self.port}_{safe_model_name}.log"
)
stdout_path.unlink(missing_ok=True)
command = [
"sglang",
"serve",
"--model-path",
self.model,
"--port",
str(self.port),
"--log-level=debug",
]
if self.extra_args.strip():
command.extend(self.extra_args.strip().split())
env = os.environ.copy()
env["SGLANG_DIFFUSION_STAGE_LOGGING"] = "1"
env["SGLANG_PERF_LOG_DIR"] = log_dir.as_posix()
# Apply custom environment variables
env.update(self.env_vars)
# TODO: unify with run_command
logger.info(f"Running command: {shlex.join(command)}")
process = subprocess.Popen(
command,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
env=env,
)
log_thread = None
stdout_fh = stdout_path.open("w", encoding="utf-8", buffering=1)
if process.stdout:
def _log_pipe(pipe: Any, file: Any) -> None:
"""Read from pipe and write to file and stdout."""
try:
with pipe:
for line in iter(pipe.readline, ""):
sys.stdout.write(line)
sys.stdout.flush()
file.write(line)
file.flush()
except Exception as e:
logger.error("Log pipe thread error: %s", e)
finally:
file.close()
logger.debug("Log pipe thread finished.")
log_thread = threading.Thread(
target=_log_pipe, args=(process.stdout, stdout_fh)
)
log_thread.daemon = True
log_thread.start()
logger.info(
"[server-test] Starting server pid=%s, model=%s, log=%s",
process.pid,
self.model,
stdout_path,
)
self._wait_for_ready(process, stdout_path)
return ServerContext(
port=self.port,
process=process,
model=self.model,
stdout_file=stdout_path,
perf_log_path=perf_log_path,
log_dir=log_dir,
_stdout_fh=stdout_fh,
_log_thread=log_thread,
)
def _wait_for_ready(self, process: subprocess.Popen, stdout_path: Path) -> None:
"""Wait for server to become ready."""
start = time.time()
ready_message = "Application startup complete."
log_period = 30
prev_log_period_count = 0
while time.time() - start < self.wait_deadline:
if process.poll() is not None:
tail = self._get_log_tail(stdout_path)
raise RuntimeError(
f"Server exited early (code {process.returncode}).\n{tail}"
)
if stdout_path.exists():
try:
content = stdout_path.read_text(encoding="utf-8", errors="ignore")
if ready_message in content:
logger.info("[server-test] Server ready")
return
except Exception as e:
logger.debug("Could not read log yet: %s", e)
elapsed = int(time.time() - start)
if (elapsed // log_period) > prev_log_period_count:
prev_log_period_count = elapsed // log_period
logger.info("[server-test] Waiting for server... elapsed=%ss", elapsed)
time.sleep(1)
tail = self._get_log_tail(stdout_path)
raise TimeoutError(f"Server not ready within {self.wait_deadline}s.\n{tail}")
@staticmethod
def _get_log_tail(path: Path, lines: int = 200) -> str:
"""Get the last N lines from a log file."""
try:
content = path.read_text(encoding="utf-8", errors="ignore")
return "\n".join(content.splitlines()[-lines:])
except Exception:
return ""
class WarmupRunner:
"""Handles warmup requests for a server."""
def __init__(
self,
port: int,
model: str,
prompt: str,
output_size: str,
output_format: str = None,
):
self.client = OpenAI(
api_key="sglang-anything",
base_url=f"http://localhost:{port}/v1",
)
self.model = model
self.prompt = prompt
self.output_size = output_size
self.output_format = output_format
def run_text_warmups(self, count: int) -> None:
"""Run text-to-image warmup requests."""
if count <= 0:
return
logger.info("[server-test] Running %s text warm-up(s)", count)
for _ in range(count):
result = self.client.images.generate(
model=self.model,
prompt=self.prompt,
n=1,
size=self.output_size,
response_format="b64_json",
)
validate_image(result.data[0].b64_json)
def run_edit_warmups(
self,
count: int,
edit_prompt: str,
image_path: Path,
) -> None:
"""Run image-edit warmup requests."""
if count <= 0:
return
if not isinstance(image_path, list):
image_path = [image_path]
for image in image_path:
if not image.exists():
logger.warning(
"[server-test] Skipping edit warmup: image missing at %s", image
)
return
logger.info("[server-test] Running %s edit warm-up(s)", count)
for _ in range(count):
images = [open(image, "rb") for image in image_path]
try:
result = self.client.images.edit(
model=self.model,
image=images,
prompt=edit_prompt,
n=1,
size=self.output_size,
response_format="b64_json",
output_format=self.output_format,
)
finally:
for img in images:
img.close()
validate_image(result.data[0].b64_json)
class PerformanceValidator:
"""Validates performance metrics against expectations."""
is_video_gen: bool = False
def __init__(
self,
scenario: ScenarioConfig,
tolerances: ToleranceConfig,
step_fractions: Sequence[float],
):
self.scenario = scenario
self.tolerances = tolerances
self.step_fractions = step_fractions
self.is_baseline_generation_mode = (
os.environ.get("SGLANG_GEN_BASELINE", "0") == "1"
)
def _assert_le(
self,
name: str,
actual: float,
expected: float,
tolerance: float,
min_abs_tolerance_ms: float = 20.0,
):
"""Assert that actual is less than or equal to expected within a tolerance.
Uses the larger of relative tolerance or absolute tolerance to prevent
flaky failures on very fast operations.
For AMD GPUs, uses 100% higher tolerance and issues warning instead of assertion.
"""
# Check if running on AMD GPU
is_amd = current_platform.is_hip()
if is_amd:
# Use 100% higher tolerance for AMD (2x the expected value)
amd_tolerance = 1.0 # 100%
upper_bound = calculate_upper_bound(
expected, amd_tolerance, min_abs_tolerance_ms
)
if actual > upper_bound:
logger.warning(
f"[AMD PERF WARNING] Validation would fail for '{name}'.\n"
f" Actual: {actual:.4f}ms\n"
f" Expected: {expected:.4f}ms\n"
f" AMD Limit: {upper_bound:.4f}ms "
f"(rel_tol: {amd_tolerance:.1%}, abs_pad: {min_abs_tolerance_ms}ms)\n"
f" Original tolerance was: {tolerance:.1%}"
)
else:
upper_bound = calculate_upper_bound(
expected, tolerance, min_abs_tolerance_ms
)
assert actual <= upper_bound, (
f"Validation failed for '{name}'.\n"
f" Actual: {actual:.4f}ms\n"
f" Expected: {expected:.4f}ms\n"
f" Limit: {upper_bound:.4f}ms "
f"(rel_tol: {tolerance:.1%}, abs_pad: {min_abs_tolerance_ms}ms)"
)
def validate(
self, perf_record: RequestPerfRecord, *args, **kwargs
) -> PerformanceSummary:
"""Validate all performance metrics and return summary."""
summary = self.collect_metrics(perf_record)
if self.is_baseline_generation_mode:
return summary
self._validate_e2e(summary)
self._validate_denoise_agg(summary)
self._validate_denoise_steps(summary)
self._validate_stages(summary)
return summary
def collect_metrics(
self,
perf_record: RequestPerfRecord,
) -> PerformanceSummary:
return PerformanceSummary.from_req_perf_record(perf_record, self.step_fractions)
def _validate_e2e(self, summary: PerformanceSummary) -> None:
"""Validate end-to-end performance."""
assert summary.e2e_ms > 0, "E2E duration missing"
self._assert_le(
"E2E Latency",
summary.e2e_ms,
self.scenario.expected_e2e_ms,
self.tolerances.e2e,
)
def _validate_denoise_agg(self, summary: PerformanceSummary) -> None:
"""Validate aggregate denoising metrics."""
assert summary.avg_denoise_ms > 0, "Denoising step timings missing"
self._assert_le(
"Average Denoise Step",
summary.avg_denoise_ms,
self.scenario.expected_avg_denoise_ms,
self.tolerances.denoise_agg,
)
self._assert_le(
"Median Denoise Step",
summary.median_denoise_ms,
self.scenario.expected_median_denoise_ms,
self.tolerances.denoise_agg,
)
def _validate_denoise_steps(self, summary: PerformanceSummary) -> None:
"""Validate individual denoising steps."""
for idx, actual in summary.sampled_steps.items():
expected = self.scenario.denoise_step_ms.get(idx)
if expected is None:
continue
# FIXME: hardcode, looser for first step
tolerance = 0.4 if idx == 0 else self.tolerances.denoise_step
self._assert_le(
f"Denoise Step {idx}",
actual,
expected,
tolerance,
)
def _validate_stages(self, summary: PerformanceSummary) -> None:
"""Validate stage-level metrics."""
assert summary.stage_metrics, "Stage metrics missing"
for stage, expected in self.scenario.stages_ms.items():
if stage == "per_frame_generation" and self.is_video_gen:
continue
actual = summary.stage_metrics.get(stage)
assert actual is not None, f"Stage {stage} timing missing"
tolerance = (
self.tolerances.denoise_stage
if stage == "DenoisingStage"
else self.tolerances.non_denoise_stage
)
self._assert_le(
f"Stage '{stage}'",
actual,
expected,
tolerance,
min_abs_tolerance_ms=120.0, # relax absolute tolerance for non-denoising stages
)
class VideoPerformanceValidator(PerformanceValidator):
"""Extended validator for video diffusion with frame-level metrics."""
is_video_gen = True
def validate(
self,
perf_record: RequestPerfRecord,
num_frames: int | None = None,
) -> PerformanceSummary:
"""Validate video metrics including frame generation rates."""
summary = super().validate(perf_record)
if num_frames and summary.e2e_ms > 0:
summary.total_frames = num_frames
summary.avg_frame_time_ms = summary.e2e_ms / num_frames
summary.frames_per_second = 1000.0 / summary.avg_frame_time_ms
if not self.is_baseline_generation_mode:
self._validate_frame_rate(summary)
return summary
def _validate_frame_rate(self, summary: PerformanceSummary) -> None:
"""Validate frame generation performance."""
expected_frame_time = self.scenario.stages_ms.get("per_frame_generation")
if expected_frame_time and summary.avg_frame_time_ms:
self._assert_le(
"Average Frame Time",
summary.avg_frame_time_ms,
expected_frame_time,
self.tolerances.denoise_stage,
)
# Registry of validators by name
VALIDATOR_REGISTRY = {
"default": PerformanceValidator,
"video": VideoPerformanceValidator,
}
def get_generate_fn(
model_path: str,
modality: str,
sampling_params: DiffusionSamplingParams,
) -> Callable[[str, Client], str]:
"""Return appropriate generation function for the case."""
# Allow override via environment variable (useful for AMD where large resolutions cause slow VAE)
output_size = os.environ.get("SGLANG_TEST_OUTPUT_SIZE", sampling_params.output_size)
n = sampling_params.num_outputs_per_prompt
def _create_and_download_video(
client,
case_id,
*,
model: str,
size: str,
prompt: str | None = None,
seconds: int | None = None,
input_reference: Any | None = None,
extra_body: dict[Any] | None = None,
) -> str:
"""
Create a video job via /v1/videos, poll until completion,
then download the binary content and validate it.
Returns request-id
"""
create_kwargs: dict[str, Any] = {
"model": model,
"size": size,
}
if prompt is not None:
create_kwargs["prompt"] = prompt
if seconds is not None:
create_kwargs["seconds"] = seconds
if input_reference is not None:
create_kwargs["input_reference"] = input_reference # triggers multipart
if extra_body is not None:
create_kwargs["extra_body"] = extra_body
job = client.videos.create(**create_kwargs) # type: ignore[attr-defined]
video_id = job.id
job_completed = False
is_baseline_generation_mode = os.environ.get("SGLANG_GEN_BASELINE", "0") == "1"
# Check if running on AMD GPU - use longer timeout
is_amd = current_platform.is_hip()
if is_baseline_generation_mode:
timeout = 3600.0
elif is_amd:
timeout = 2400.0 # 40 minutes for AMD
else:
timeout = 1200.0
deadline = time.time() + timeout
while True:
page = client.videos.list() # type: ignore[attr-defined]
item = next((v for v in page.data if v.id == video_id), None)
if item and getattr(item, "status", None) == "completed":
job_completed = True
break
if time.time() > deadline:
break
time.sleep(1)
if not job_completed:
if is_baseline_generation_mode:
logger.warning(
f"{case_id}: video job {video_id} timed out during baseline generation. "
"Attempting to collect performance data anyway."
)
return video_id
if is_amd:
logger.warning(
f"[AMD TIMEOUT WARNING] {case_id}: video job {video_id} did not complete "
f"within {timeout}s timeout. This may indicate performance issues on AMD."
)
pytest.skip(
f"{case_id}: video job timed out on AMD after {timeout}s - skipping"
)
pytest.fail(f"{case_id}: video job {video_id} did not complete in time")
# download video
resp = client.videos.download_content(video_id=video_id) # type: ignore[attr-defined]
content = resp.read()
validate_openai_video(content)
expected_filename = f"{video_id}.mp4"
tmp_path = expected_filename
with open(tmp_path, "wb") as f:
f.write(content)
# Validate output file
expected_width, expected_height = parse_dimensions(size)
validate_video_file(
tmp_path, expected_filename, expected_width, expected_height
)
upload_file_to_slack(
case_id=case_id,
model=model_path,
prompt=sampling_params.prompt,
file_path=tmp_path,
origin_file_path=sampling_params.image_path,
)
os.remove(tmp_path)
return video_id
video_seconds = sampling_params.seconds or 4
def generate_image(case_id, client) -> str:
"""T2I: Text to Image generation."""
if not sampling_params.prompt:
pytest.skip(f"{id}: no text prompt configured")
# Request parameters that affect output format
req_output_format = None # Not specified in current request
req_background = None # Not specified in current request
# Build extra_body for optional features
extra_body = {}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
response = client.images.with_raw_response.generate(
model=model_path,
prompt=sampling_params.prompt,
n=n,
size=output_size,
response_format="b64_json",
extra_body=extra_body if extra_body else None,
)
result = response.parse()
validate_image(result.data[0].b64_json)
rid = result.id
img_data = base64.b64decode(result.data[0].b64_json)
# Infer expected format from request parameters
expected_ext = get_expected_image_format(req_output_format, req_background)
expected_filename = f"{result.created}.{expected_ext}"
tmp_path = expected_filename
with open(tmp_path, "wb") as f:
f.write(img_data)
# Validate output file
expected_width, expected_height = parse_dimensions(output_size)
validate_image_file(
tmp_path,
expected_filename,
expected_width,
expected_height,
output_format=req_output_format,
background=req_background,
)
upload_file_to_slack(
case_id=case_id,
model=model_path,
prompt=sampling_params.prompt,
file_path=tmp_path,
)
os.remove(tmp_path)
return rid
def generate_image_edit(case_id, client) -> str:
"""TI2I: Text + Image ? Image edit."""
if not sampling_params.prompt or not sampling_params.image_path:
pytest.skip(f"{id}: no edit config")
image_paths = sampling_params.image_path
if not isinstance(image_paths, list):
image_paths = [image_paths]
new_image_paths = []
for image_path in image_paths:
if is_image_url(image_path):
new_image_paths.append(download_image_from_url(str(image_path)))
else:
new_image_paths.append(Path(image_path))
if not image_path.exists():
pytest.skip(f"{id}: file missing: {image_path}")
image_paths = new_image_paths
# Request parameters that affect output format
req_output_format = (
sampling_params.output_format
) # Not specified in current request
req_background = None # Not specified in current request
# Build extra_body for optional features
extra_body = {"num_frames": sampling_params.num_frames}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
images = [open(image_path, "rb") for image_path in image_paths]
try:
response = client.images.with_raw_response.edit(
model=model_path,
image=images,
prompt=sampling_params.prompt,
n=n,
size=output_size,
response_format="b64_json",
output_format=req_output_format,
extra_body=extra_body,
)
finally:
for img in images:
img.close()
result = response.parse()
validate_image(result.data[0].b64_json)
img_data = base64.b64decode(result.data[0].b64_json)
rid = result.id
# Infer expected format from request parameters
expected_ext = get_expected_image_format(req_output_format, req_background)
expected_filename = f"{rid}.{expected_ext}"
tmp_path = expected_filename
with open(tmp_path, "wb") as f:
f.write(img_data)
# Validate output file
expected_width, expected_height = parse_dimensions(output_size)
validate_image_file(
tmp_path,
expected_filename,
expected_width,
expected_height,
output_format=req_output_format,
background=req_background,
)
upload_file_to_slack(
case_id=case_id,
model=model_path,
prompt=sampling_params.prompt,
file_path=tmp_path,
origin_file_path=sampling_params.image_path,
)
os.remove(tmp_path)
return rid
def generate_image_edit_url(case_id, client) -> str:
"""TI2I: Text + Image ? Image edit using direct URL transfer (no pre-download)."""
if not sampling_params.prompt or not sampling_params.image_path:
pytest.skip(f"{id}: no edit config")
# Handle both single URL and list of URLs
image_urls = sampling_params.image_path
if not isinstance(image_urls, list):
image_urls = [image_urls]
# Validate all URLs
for url in image_urls:
if not is_image_url(url):
pytest.skip(
f"{id}: image_path must be a URL for URL direct test: {url}"
)
# Request parameters that affect output format
req_output_format = (
sampling_params.output_format
) # Not specified in current request
req_background = None # Not specified in current request
response = client.images.with_raw_response.edit(
model=model_path,
prompt=sampling_params.prompt,
image=[], # Only for OpenAI verification
n=n,
size=sampling_params.output_size,
response_format="b64_json",
output_format=req_output_format,
extra_body={"url": image_urls, "num_frames": sampling_params.num_frames},
)
result = response.parse()
rid = result.id
validate_image(result.data[0].b64_json)
# Save and upload result for verification
img_data = base64.b64decode(result.data[0].b64_json)
# Infer expected format from request parameters
expected_ext = get_expected_image_format(req_output_format, req_background)
expected_filename = f"{rid}.{expected_ext}"
tmp_path = expected_filename
with open(tmp_path, "wb") as f:
f.write(img_data)
# Validate output file
expected_width, expected_height = parse_dimensions(sampling_params.output_size)
validate_image_file(
tmp_path,
expected_filename,
expected_width,
expected_height,
output_format=req_output_format,
background=req_background,
)
upload_file_to_slack(
case_id=case_id,
model=model_path,
prompt=sampling_params.prompt,
file_path=tmp_path,
origin_file_path=str(sampling_params.image_path),
)
os.remove(tmp_path)
return rid
def generate_video(case_id, client) -> str:
"""T2V: Text ? Video."""
if not sampling_params.prompt:
pytest.skip(f"{id}: no text prompt configured")
# Build extra_body for optional features
extra_body = {}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
return _create_and_download_video(
client,
case_id,
model=model_path,
prompt=sampling_params.prompt,
size=output_size,
seconds=video_seconds,
extra_body=extra_body if extra_body else None,
)
def generate_image_to_video(case_id, client) -> str:
"""I2V: Image ? Video (optional prompt)."""
if not sampling_params.image_path:
pytest.skip(f"{id}: no input image configured")
if is_image_url(sampling_params.image_path):
image_path = download_image_from_url(str(sampling_params.image_path))
else:
image_path = Path(sampling_params.image_path)
if not image_path.exists():
pytest.skip(f"{id}: file missing: {image_path}")
# Build extra_body for optional features
extra_body = {}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
with image_path.open("rb") as fh:
return _create_and_download_video(
client,
case_id,
model=model_path,
prompt=sampling_params.prompt,
size=output_size,
seconds=video_seconds,
input_reference=fh,
extra_body=extra_body if extra_body else None,
)
def generate_text_url_image_to_video(case_id, client) -> str:
if not sampling_params.prompt or not sampling_params.image_path:
pytest.skip(f"{id}: no edit config")
# Build extra_body for optional features
extra_body = {"reference_url": sampling_params.image_path}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
return _create_and_download_video(
client,
case_id,
model=model_path,
prompt=sampling_params.prompt,
size=sampling_params.output_size,
seconds=video_seconds,
extra_body={
"reference_url": sampling_params.image_path,
"fps": sampling_params.fps,
"num_frames": sampling_params.num_frames,
},
)
def generate_text_image_to_video(case_id, client) -> str:
"""TI2V: Text + Image ? Video."""
if not sampling_params.prompt or not sampling_params.image_path:
pytest.skip(f"{id}: no edit config")
if is_image_url(sampling_params.image_path):
image_path = download_image_from_url(str(sampling_params.image_path))
else:
image_path = Path(sampling_params.image_path)
if not image_path.exists():
pytest.skip(f"{id}: file missing: {image_path}")
# Build extra_body for optional features
extra_body = {}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
with image_path.open("rb") as fh:
return _create_and_download_video(
client,
case_id,
model=model_path,
prompt=sampling_params.prompt,
size=output_size,
seconds=video_seconds,
input_reference=fh,
extra_body={
"fps": sampling_params.fps,
"num_frames": sampling_params.num_frames,
},
)
if modality == "video":
if sampling_params.image_path and sampling_params.prompt:
if getattr(sampling_params, "direct_url_test", False):
fn = generate_text_url_image_to_video
else:
fn = generate_text_image_to_video
elif sampling_params.image_path:
fn = generate_image_to_video
else:
fn = generate_video
elif sampling_params.prompt and sampling_params.image_path:
if getattr(sampling_params, "direct_url_test", False):
fn = generate_image_edit_url
else:
fn = generate_image_edit
else:
fn = generate_image
return fn