Spaces:
Sleeping
Sleeping
| """Image editing tool using fal.ai API.""" | |
| import os | |
| import json | |
| import time | |
| import tempfile | |
| from typing import List, Union | |
| from concurrent.futures import ThreadPoolExecutor, as_completed | |
| import fal_client | |
| import requests | |
| from pydantic import BaseModel, Field, field_validator | |
| from langchain_core.tools import tool | |
| # Configure fal.ai | |
| fal_client.api_key = os.getenv("FAL_KEY") | |
| EDIT_MODEL = "fal-ai/bytedance/seedream/v4.5/edit" | |
| class ImageEditInput(BaseModel): | |
| """Input schema for image editing.""" | |
| image_paths: Union[str, List[str]] = Field( | |
| description="File path(s) of images to edit. Can be a single path, comma-separated paths, or a list of paths." | |
| ) | |
| prompt: str = Field( | |
| description="Text description of how to edit the image(s). Be specific about the changes you want." | |
| ) | |
| def normalize_paths(cls, v): | |
| """Convert list to comma-separated string for consistent handling.""" | |
| if isinstance(v, list): | |
| return ", ".join(str(p) for p in v) | |
| return v | |
| def _upload_image_to_fal(image_path: str) -> str: | |
| """Upload a local image to fal.ai CDN and return the URL.""" | |
| with open(image_path, "rb") as f: | |
| image_data = f.read() | |
| # Determine content type | |
| if image_path.lower().endswith(".png"): | |
| content_type = "image/png" | |
| elif image_path.lower().endswith((".jpg", ".jpeg")): | |
| content_type = "image/jpeg" | |
| elif image_path.lower().endswith(".webp"): | |
| content_type = "image/webp" | |
| else: | |
| content_type = "image/png" | |
| return fal_client.upload(image_data, content_type) | |
| def _edit_single_image(image_path: str, prompt: str) -> str: | |
| """Edit a single image using fal.ai API.""" | |
| # Upload image to fal.ai CDN | |
| image_url = _upload_image_to_fal(image_path) | |
| # API expects image_urls as a list (matching smolagents implementation) | |
| args = { | |
| "image_urls": [image_url], | |
| "prompt": prompt, | |
| "width": 1024, | |
| "height": 1024, | |
| } | |
| result = fal_client.submit(EDIT_MODEL, arguments=args).get() | |
| if not result or "images" not in result or not result["images"]: | |
| raise ValueError("No edited image returned by fal.ai") | |
| output_url = result["images"][0]["url"] | |
| # Download the edited image | |
| response = requests.get(output_url, timeout=60) | |
| response.raise_for_status() | |
| temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False) | |
| temp_file.write(response.content) | |
| temp_file.close() | |
| return temp_file.name | |
| def edit_images(image_paths: str, prompt: str) -> str: | |
| """Edit existing images based on a text prompt. | |
| Use this tool when the user wants to modify, transform, or change existing images. | |
| The user must provide the file path(s) of images to edit. | |
| Args: | |
| image_paths: Comma-separated file paths of images to edit | |
| prompt: Description of how to edit the image(s) | |
| Returns: | |
| String containing the file paths of edited images, separated by newlines | |
| """ | |
| # Parse paths | |
| paths = [p.strip() for p in image_paths.split(",") if p.strip()] | |
| # Debug logging | |
| debug_info = { | |
| "image_paths": image_paths, | |
| "parsed_paths": paths, | |
| "prompt": prompt[:100] + "..." if len(prompt) > 100 else prompt, | |
| "num_images": len(paths), | |
| } | |
| print(f"\n{'='*80}") | |
| print(f"DEBUG - Image Edit Tool Called:") | |
| print(json.dumps(debug_info, indent=2)) | |
| print(f"{'='*80}\n") | |
| if not paths: | |
| print(f"DEBUG - ERROR: No valid image paths provided") | |
| return "Error: No valid image paths provided" | |
| # Validate paths exist | |
| valid_paths = [] | |
| for path in paths: | |
| if os.path.exists(path): | |
| valid_paths.append(path) | |
| print(f"DEBUG - Valid path: {path}") | |
| else: | |
| print(f"DEBUG - ERROR: File not found: {path}") | |
| return f"Error: Image file not found: {path}" | |
| print(f"DEBUG - {len(valid_paths)} valid paths to edit") | |
| edited_paths = [] | |
| errors = [] | |
| st = time.time() | |
| if len(valid_paths) == 1: | |
| try: | |
| print(f"DEBUG - Editing single image: {valid_paths[0]}") | |
| edited_path = _edit_single_image(valid_paths[0], prompt) | |
| edited_paths.append(edited_path) | |
| print(f"DEBUG - Single image edited: {edited_path}") | |
| except Exception as e: | |
| print(f"DEBUG - Single image edit EXCEPTION: {type(e).__name__}: {e}") | |
| errors.append(f"{valid_paths[0]}: {str(e)}") | |
| else: | |
| # Parallel editing for multiple images | |
| print(f"DEBUG - Parallel editing {len(valid_paths)} images...") | |
| with ThreadPoolExecutor(max_workers=len(valid_paths)) as executor: | |
| futures = { | |
| executor.submit(_edit_single_image, path, prompt): path | |
| for path in valid_paths | |
| } | |
| for future in as_completed(futures): | |
| original_path = futures[future] | |
| try: | |
| edited_path = future.result() | |
| edited_paths.append(edited_path) | |
| print(f"DEBUG - Edited: {original_path} -> {edited_path}") | |
| except Exception as e: | |
| print(f"DEBUG - Edit EXCEPTION for {original_path}: {type(e).__name__}: {e}") | |
| errors.append(f"{original_path}: {str(e)}") | |
| et = time.time() | |
| print(f"DEBUG - Image editing completed in {et - st:.2f}s") | |
| if not edited_paths: | |
| print(f"DEBUG - All edits failed: {errors}") | |
| return f"Failed to edit images: {'; '.join(errors)}" | |
| print(f"DEBUG - Successfully edited {len(edited_paths)} image(s)") | |
| for path in edited_paths: | |
| print(f"DEBUG - Edited image path: {path}") | |
| result = f"Successfully edited {len(edited_paths)} image(s):\n" | |
| for path in edited_paths: | |
| result += f"- {path}\n" | |
| if errors: | |
| result += f"\nSome edits failed: {'; '.join(errors)}" | |
| print(f"DEBUG - Image edit result:\n{result}") | |
| return result | |