""" Unified tool router — maps task intent to the best HF Space and calls it. All spaces are already MCP-wired on bovi1's account. """ from gradio_client import Client, handle_file from PIL import Image import tempfile, os, requests # ── Tool registry ────────────────────────────────────────────────────────────── # Each entry: space_id, description, input_schema, fn_index or api_name TOOLS = { # Image generation "flux": {"space": "black-forest-labs/FLUX.2-klein-9B", "task": "text_to_image", "desc": "High-quality text-to-image (FLUX)"}, "z_turbo": {"space": "mcp-tools/Z-Image-Turbo", "task": "text_to_image", "desc": "Fast text-to-image (Z-Turbo)"}, "firered": {"space": "prithivMLmods/FireRed-Image-Edit-1.0-Fast", "task": "image_edit", "desc": "Fast image editing (FireRed)"}, "qwen_edit": {"space": "Onise/Qwen-Image-Edit-2509-LoRAs-Fast2", "task": "image_edit", "desc": "Qwen image editing with LoRAs"}, "qwen_pose": {"space": "linoyts/Qwen-Image-Edit-2511-AnyPose", "task": "image_edit", "desc": "Pose-guided image editing"}, "flux_lora": {"space": "M3st3rJ4k3l/FLUX.2-Klein-Multi-LoRA", "task": "text_to_image", "desc": "FLUX with multi-LoRA support"}, # Video generation "wan2_fast": {"space": "zerogpu-aoti/wan2-2-fp8da-aoti-faster", "task": "text_to_video", "desc": "Fastest WAN2.2 text-to-video"}, "wan2_img": {"space": "EldMans/wan2.2_14b_i2v_480p_lightning_nsfw_diffusers", "task": "image_to_video", "desc": "WAN2.2 image-to-video"}, "ltx": {"space": "alexnasa/ltx-2-TURBO", "task": "text_to_video", "desc": "LTX-2 fast video gen"}, # Face "face_parse": {"space": "leonelhs/faceparser", "task": "face_parse", "desc": "Detailed face region parsing"}, "face_detect": {"space": "KaiSKX/vit_fACE_detection", "task": "face_detect", "desc": "ViT face detection"}, "face_anal": {"space": "leonelhs/FaceAnalysis", "task": "face_analyze", "desc": "Age/gender/emotion analysis"}, # Pose "vitpose": {"space": "hysts/ViTPose-transformers", "task": "pose_estimate", "desc": "Body pose estimation"}, # Segmentation / background "bg_remove": {"space": "not-lain/background-removal", "task": "background_remove","desc": "Clean background removal"}, "sam3": {"space": "prithivMLmods/SAM3-Demo", "task": "segment", "desc": "SAM3 segmentation"}, # 3D "trellis": {"space": "trellis-community/TRELLIS", "task": "image_to_3d", "desc": "Image to 3D mesh (TRELLIS)"}, # Upscale "upscale": {"space": "prithivMLmods/PiD-Image-Upscaler", "task": "upscale", "desc": "Image upscaling"}, # Audio "tts": {"space": "ResembleAI/Chatterbox-Multilingual-TTS-V3", "task": "tts", "desc": "Multilingual TTS (Chatterbox)"}, } # ── Intent → best tool mapping ───────────────────────────────────────────────── INTENT_MAP = { "generate image": "flux", "text to image": "flux", "fast image": "z_turbo", "edit image": "qwen_edit", "edit pose": "qwen_pose", "change pose": "qwen_pose", "image editing": "firered", "generate video": "wan2_fast", "text to video": "wan2_fast", "image to video": "wan2_img", "fast video": "ltx", "parse face": "face_parse", "detect face": "face_detect", "face analysis": "face_anal", "analyze face": "face_anal", "pose estimation": "vitpose", "estimate pose": "vitpose", "remove background": "bg_remove", "background removal": "bg_remove", "segment": "sam3", "3d": "trellis", "image to 3d": "trellis", "upscale": "upscale", "enhance": "upscale", "tts": "tts", "text to speech": "tts", "speak": "tts", } def resolve_tool(task: str, tool_override: str = None) -> dict: """Pick the best tool for a given task string.""" if tool_override and tool_override in TOOLS: return {"key": tool_override, **TOOLS[tool_override]} task_lower = task.lower() for phrase, key in INTENT_MAP.items(): if phrase in task_lower: return {"key": key, **TOOLS[key]} return None def call_tool(tool_key: str, inputs: dict, hf_token: str = None) -> dict: """ Call a space via Gradio client. inputs: dict with keys like prompt, image, audio depending on task """ tool = TOOLS.get(tool_key) if not tool: return {"error": f"Unknown tool: {tool_key}"} try: kwargs = {"src": tool["space"]} if hf_token: kwargs["hf_token"] = hf_token client = Client(**kwargs) task = tool["task"] prompt = inputs.get("prompt", "") image = inputs.get("image") # PIL Image or file path # Save image to temp file if provided img_path = None if image is not None: if isinstance(image, Image.Image): tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False) image.save(tmp.name) img_path = tmp.name else: img_path = image if task == "text_to_image": result = client.predict(prompt, api_name="/predict") elif task == "image_edit": result = client.predict(handle_file(img_path), prompt, api_name="/predict") elif task == "text_to_video": result = client.predict(prompt, api_name="/predict") elif task == "image_to_video": result = client.predict(handle_file(img_path), prompt, api_name="/predict") elif task in ("face_parse", "face_detect", "face_analyze"): result = client.predict(handle_file(img_path), api_name="/predict") elif task == "pose_estimate": result = client.predict(handle_file(img_path), api_name="/predict") elif task == "background_remove": result = client.predict(handle_file(img_path), api_name="/predict") elif task == "segment": result = client.predict(handle_file(img_path), api_name="/predict") elif task == "image_to_3d": result = client.predict(handle_file(img_path), api_name="/predict") elif task == "upscale": result = client.predict(handle_file(img_path), api_name="/predict") elif task == "tts": text = inputs.get("text", prompt) lang = inputs.get("language", "en") result = client.predict(text, lang, api_name="/predict") else: result = client.predict(api_name="/predict") if img_path and img_path != image: os.unlink(img_path) return {"tool": tool_key, "space": tool["space"], "result": result} except Exception as e: return {"tool": tool_key, "space": tool["space"], "error": str(e)} def list_tools() -> list[dict]: return [ {"key": k, "space": v["space"], "task": v["task"], "desc": v["desc"]} for k, v in TOOLS.items() ]