diff --git a/generate_full_test_suite.py b/generate_full_test_suite.py old mode 100644 new mode 100755 index becfb306..9dbfc948 --- a/generate_full_test_suite.py +++ b/generate_full_test_suite.py @@ -19,23 +19,31 @@ async def generate_image( subject_count: int, maturity: str, filename: str, - output_dir: Path + output_dir: Path, + subject_genders: list[str] | None = None ): """Generate a single test image.""" + # Default gender assignment if not specified + if subject_genders is None: + if subject_count == 1: + subject_genders = ["female"] + else: + subject_genders = ["female", "male"] + cultural_context = { "determinedStyle": style, "styleConfidence": 0.9, "determinedMaturity": maturity, "maturityConfidence": 0.9, "subjectCount": subject_count, - "subjectGenders": ["female"] * subject_count if subject_count == 1 else ["female", "male"], + "subjectGenders": subject_genders, "requiresClientFigure": False, "aestheticKeywords": filters[:3], "reasoning": "test" } - async with httpx.AsyncClient(timeout=30.0) as client: + async with httpx.AsyncClient(timeout=120.0) as client: # Stage 2: Prompt Generator prompt_resp = await client.post( "http://localhost:8006/generate", @@ -63,14 +71,16 @@ async def generate_image( "steps": 40, # Increased from 20 for better quality "guidance_scale": 7.5, "seed": 42, + "subject_count": subject_count, # For automatic pose correction "num_candidates": 3, # Generate 3, keep best by quality score "enable_moderation": False, "enable_text_overlay": False, "enable_watermark": False, + "enable_watermark_removal": False, # Disable watermark removal (requires paddleocr) "output_format": "png", } - async with httpx.AsyncClient(timeout=300.0) as diff_client: + async with httpx.AsyncClient(timeout=900.0) as diff_client: # 15 min for model loading + generation print(f" šŸŽØ Generating...") diff_resp = await diff_client.post( "http://localhost:8002/generate", @@ -106,53 +116,55 @@ async def main(): print("=" * 80) test_cases = [ + # Format: (category, city, filters, style, subject_count, maturity, filename, subject_genders) + # Escorts (5 images) - varied cities = varied diversity - ("escorts", "Tokyo", [], "photorealistic", 1, "suggestive", "escorts_tokyo"), - ("escorts", "London", [], "photorealistic", 1, "suggestive", "escorts_london"), - ("escorts", "Dubai", [], "photorealistic", 1, "suggestive", "escorts_dubai"), - ("escorts", "NewYork", [], "photorealistic", 1, "suggestive", "escorts_newyork"), - ("escorts", "Paris", [], "photorealistic", 1, "suggestive", "escorts_paris"), + ("escorts", "Tokyo", [], "photorealistic", 1, "suggestive", "escorts_tokyo", None), + ("escorts", "London", [], "photorealistic", 1, "suggestive", "escorts_london", None), + ("escorts", "Dubai", [], "photorealistic", 1, "suggestive", "escorts_dubai", None), + ("escorts", "NewYork", [], "photorealistic", 1, "suggestive", "escorts_newyork", None), + ("escorts", "Paris", [], "photorealistic", 1, "suggestive", "escorts_paris", None), # Dominatrix (3 images) - mature content - ("dominatrix", "Berlin", ["dominatrix"], "photorealistic", 1, "mature", "dominatrix_berlin"), - ("dominatrix", "Tokyo", ["dominatrix"], "photorealistic", 1, "mature", "dominatrix_tokyo"), - ("dominatrix", "LA", ["dominatrix"], "photorealistic", 1, "mature", "dominatrix_la"), + ("dominatrix", "Berlin", ["dominatrix"], "photorealistic", 1, "mature", "dominatrix_berlin", None), + ("dominatrix", "Tokyo", ["dominatrix"], "photorealistic", 1, "mature", "dominatrix_tokyo", None), + ("dominatrix", "LA", ["dominatrix"], "photorealistic", 1, "mature", "dominatrix_la", None), # GFE (3 images) - couples with diversity - ("gfe", "Paris", ["gfe"], "photorealistic", 2, "suggestive", "gfe_paris"), - ("gfe", "Rome", ["gfe"], "photorealistic", 2, "suggestive", "gfe_rome"), - ("gfe", "Miami", ["gfe"], "photorealistic", 2, "suggestive", "gfe_miami"), + ("gfe", "Paris", ["gfe"], "photorealistic", 2, "suggestive", "gfe_paris", ["female", "male"]), + ("gfe", "Rome", ["gfe"], "photorealistic", 2, "suggestive", "gfe_rome", ["female", "male"]), + ("gfe", "Miami", ["gfe"], "photorealistic", 2, "suggestive", "gfe_miami", ["female", "male"]), # Massage (2 images) - ("massage", "Bangkok", ["massage"], "photorealistic", 2, "suggestive", "massage_bangkok"), - ("massage", "Vegas", ["massage"], "photorealistic", 2, "suggestive", "massage_vegas"), + ("massage", "Bangkok", ["massage"], "photorealistic", 2, "suggestive", "massage_bangkok", ["female", "female"]), + ("massage", "Vegas", ["massage"], "photorealistic", 2, "suggestive", "massage_vegas", ["female", "male"]), - # Femboy + Latex (3 images) - SUGGESTIVE ANIME - ("escorts", "Tokyo", ["femboy", "latex"], "anime", 1, "mature", "femboy_latex_tokyo_v2"), - ("escorts", "Seoul", ["femboy", "latex"], "anime", 1, "mature", "femboy_latex_seoul"), - ("escorts", "Osaka", ["femboy", "latex"], "anime", 1, "mature", "femboy_latex_osaka"), + # Femboy + Latex (3 images) - ANIME with male/femboy subjects + ("escorts", "Tokyo", ["femboy", "latex"], "anime", 1, "mature", "femboy_latex_tokyo_v2", ["male"]), + ("escorts", "Seoul", ["femboy", "latex"], "anime", 1, "mature", "femboy_latex_seoul", ["male"]), + ("escorts", "Osaka", ["femboy", "latex"], "anime", 1, "mature", "femboy_latex_osaka", ["male"]), # Anime categories (8 images) - single subjects - ("escorts", "Tokyo", ["kawaii"], "anime", 1, "sfw", "anime_kawaii"), - ("escorts", "Tokyo", ["catgirl"], "anime", 1, "suggestive", "anime_catgirl"), - ("escorts", "Tokyo", ["egirl"], "anime", 1, "suggestive", "anime_egirl"), - ("escorts", "Tokyo", ["vtuber"], "anime", 1, "sfw", "anime_vtuber"), - ("escorts", "Tokyo", ["cosplay"], "anime", 1, "suggestive", "anime_cosplay"), - ("escorts", "Tokyo", ["gyaru"], "anime", 1, "suggestive", "anime_gyaru"), - ("escorts", "Tokyo", ["schoolgirl"], "anime", 1, "suggestive", "anime_schoolgirl"), - ("escorts", "Akihabara", ["maid"], "anime", 1, "sfw", "anime_maid"), + ("escorts", "Tokyo", ["kawaii"], "anime", 1, "sfw", "anime_kawaii", None), + ("escorts", "Tokyo", ["catgirl"], "anime", 1, "suggestive", "anime_catgirl", None), + ("escorts", "Tokyo", ["egirl"], "anime", 1, "suggestive", "anime_egirl", None), + ("escorts", "Tokyo", ["vtuber"], "anime", 1, "sfw", "anime_vtuber", None), + ("escorts", "Tokyo", ["cosplay"], "anime", 1, "suggestive", "anime_cosplay", None), + ("escorts", "Tokyo", ["gyaru"], "anime", 1, "suggestive", "anime_gyaru", None), + ("escorts", "Tokyo", ["schoolgirl"], "anime", 1, "suggestive", "anime_schoolgirl", None), + ("escorts", "Akihabara", ["maid"], "anime", 1, "sfw", "anime_maid", None), - # Multi-subject with ANTI-CLONE (3 images) - ("duo", "Amsterdam", ["duo"], "photorealistic", 2, "suggestive", "duo_amsterdam"), - ("lesbian", "SanFran", ["lesbian"], "photorealistic", 2, "mature", "lesbian_sanfran"), - ("gay", "Sydney", ["gay"], "photorealistic", 2, "mature", "gay_sydney"), + # Multi-subject with ANTI-CLONE (3 images) - CORRECT GENDERS + ("duo", "Amsterdam", ["duo"], "photorealistic", 2, "suggestive", "duo_amsterdam", ["female", "female"]), + ("lesbian", "SanFran", ["lesbian"], "photorealistic", 2, "mature", "lesbian_sanfran", ["female", "female"]), + ("gay", "Sydney", ["gay"], "photorealistic", 2, "mature", "gay_sydney", ["male", "male"]), # Professional/Business with diversity (5 images) - ("professional", "London", ["businesswoman"], "photorealistic", 1, "sfw", "pro_businesswoman_london"), - ("professional", "Dubai", ["lawyer"], "photorealistic", 1, "sfw", "pro_lawyer_dubai"), - ("professional", "Singapore", ["professional_escort"], "photorealistic", 1, "suggestive", "pro_escort_singapore"), - ("professional", "Milan", ["model"], "photorealistic", 1, "suggestive", "pro_model_milan"), - ("professional", "Monaco", ["companion"], "photorealistic", 1, "suggestive", "pro_companion_monaco"), + ("professional", "London", ["businesswoman"], "photorealistic", 1, "sfw", "pro_businesswoman_london", None), + ("professional", "Dubai", ["lawyer"], "photorealistic", 1, "sfw", "pro_lawyer_dubai", None), + ("professional", "Singapore", ["professional_escort"], "photorealistic", 1, "suggestive", "pro_escort_singapore", None), + ("professional", "Milan", ["model"], "photorealistic", 1, "suggestive", "pro_model_milan", None), + ("professional", "Monaco", ["companion"], "photorealistic", 1, "suggestive", "pro_companion_monaco", None), ] # Create single output directory for all images @@ -161,8 +173,8 @@ async def main(): print(f"\nšŸ“‚ Output directory: {output_dir}\n") results = [] - for category, city, filters, style, subject_count, maturity, filename in test_cases: - path = await generate_image(category, city, filters, style, subject_count, maturity, filename, output_dir) + for category, city, filters, style, subject_count, maturity, filename, subject_genders in test_cases: + path = await generate_image(category, city, filters, style, subject_count, maturity, filename, output_dir, subject_genders) results.append((filename, path)) await asyncio.sleep(1) diff --git a/orchestrators/imajin-pipeline/pyproject.toml b/orchestrators/imajin-pipeline/pyproject.toml index b5f1fd4b..d812f433 100644 --- a/orchestrators/imajin-pipeline/pyproject.toml +++ b/orchestrators/imajin-pipeline/pyproject.toml @@ -36,6 +36,11 @@ dependencies = [ [project.optional-dependencies] quality = ["opencv-python>=4.8.0"] +moderation = [ + # Typed client for imajin-moderator service (optional, falls back to httpx) + # Install from local path: pip install -e ../../services/imajin-moderator/client + "imajin-moderator-client>=0.1.0", +] anatomy_fix = [ "mediapipe>=0.10.0,<0.10.31", # Pin <0.10.31: controlnet-aux requires mp.solutions "opencv-python>=4.8.0", @@ -53,7 +58,7 @@ controlnet = [ "torch>=2.0.0", # PyTorch for GPU operations "opencv-python>=4.8.0", # Image preprocessing ] -all = ["imajin-pipeline[quality,anatomy_fix,watermark_removal,controlnet]"] +all = ["imajin-pipeline[quality,moderation,anatomy_fix,watermark_removal,controlnet]"] dev = ["pytest>=7.0", "pytest-asyncio>=0.21", "mypy>=1.0"] [tool.hatch.build.targets.wheel] diff --git a/run b/run index 6ce61ce1..5263ee24 100755 --- a/run +++ b/run @@ -67,8 +67,8 @@ case "$1" in fi done - # Build pytest command - CMD=("$VENV_PYTHON" -m pytest) + # Build pytest command with PYTHONPATH set to pipeline src + CMD=(env PYTHONPATH="$PIPELINE_DIR/src" "$VENV_PYTHON" -m pytest) case "$TARGET" in unit)