Spaces:
Sleeping
Sleeping
| """Image generation tool using fal.ai APIs.""" | |
| import os | |
| import json | |
| import time | |
| import tempfile | |
| from io import BytesIO | |
| from typing import Optional, Literal, List | |
| import fal_client | |
| import requests | |
| from pydantic import BaseModel, Field | |
| from langchain_core.tools import tool | |
| from PIL import Image | |
| # Configure fal.ai | |
| fal_client.api_key = os.getenv("FAL_KEY") | |
| # Model mapping for pro mode | |
| PRO_MODELS = { | |
| "seedream4": "fal-ai/bytedance/seedream/v4/text-to-image", | |
| "imagen4_ultra": "fal-ai/imagen4/preview/ultra", | |
| "flux-krea": "fal-ai/flux/krea", | |
| } | |
| class ImageGenerationInput(BaseModel): | |
| """Input schema for image generation.""" | |
| prompt: str = Field(description="Detailed text description of the image to generate") | |
| mode: Literal["fast", "pro"] = Field( | |
| default="fast", | |
| description="Generation mode: 'fast' uses z-image turbo for quick results, 'pro' uses premium fal.ai models" | |
| ) | |
| model: str = Field( | |
| default="flux-krea", | |
| description="Model for pro mode: 'seedream4', 'imagen4_ultra', 'flux-krea'. Default: flux-krea" | |
| ) | |
| width: int = Field(default=720, description="Image width in pixels (256-2048)") | |
| height: int = Field(default=720, description="Image height in pixels (256-2048)") | |
| num_images: int = Field(default=4, description="Number of images to generate (1-4)") | |
| negative_prompt: Optional[str] = Field( | |
| default=None, | |
| description="Things to avoid in the generated image" | |
| ) | |
| def _download_and_save_images(image_results: list) -> List[str]: | |
| """Download images from URLs and save to temp files.""" | |
| saved_paths = [] | |
| for img_data in image_results: | |
| try: | |
| img_url = img_data.get("url") if isinstance(img_data, dict) else str(img_data) | |
| if not img_url: | |
| continue | |
| response = requests.get(img_url, timeout=30) | |
| response.raise_for_status() | |
| img = Image.open(BytesIO(response.content)) | |
| temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png") | |
| img.save(temp_file.name) | |
| saved_paths.append(temp_file.name) | |
| except Exception as e: | |
| print(f"DEBUG - Error downloading image: {e}") | |
| return saved_paths | |
| def _generate_fast(prompt: str, width: int, height: int, num_images: int) -> List[str]: | |
| """Generate images using z-image turbo (fast mode).""" | |
| result = fal_client.subscribe( | |
| "fal-ai/z-image/turbo", | |
| arguments={ | |
| "prompt": prompt, | |
| "image_size": {"width": width, "height": height}, | |
| "num_images": num_images, | |
| "acceleration": "high", | |
| "num_inference_steps": 8, | |
| }, | |
| with_logs=True, | |
| ) | |
| if "images" in result and result["images"]: | |
| return _download_and_save_images(result["images"]) | |
| raise ValueError("No images generated by z-image turbo") | |
| def _generate_pro(prompt: str, model_id: str, width: int, height: int, num_images: int) -> List[str]: | |
| """Generate images using premium fal.ai models (pro mode).""" | |
| args = { | |
| "prompt": prompt, | |
| "image_size": {"width": width, "height": height}, | |
| "num_images": num_images, | |
| } | |
| handler = fal_client.submit(model_id, arguments=args) | |
| result = handler.get() | |
| if "images" in result and result["images"]: | |
| return _download_and_save_images(result["images"]) | |
| raise ValueError(f"No images generated by {model_id}") | |
| def generate_images( | |
| prompt: str, | |
| mode: Literal["fast", "pro"] = "fast", | |
| model: str = "flux-krea", | |
| width: int = 720, | |
| height: int = 720, | |
| num_images: int = 4, | |
| negative_prompt: Optional[str] = None | |
| ) -> str: | |
| """Generate images from a text prompt using AI models. | |
| Use this tool when the user wants to create new images from scratch based on a description. | |
| Returns file paths to the generated images. | |
| Args: | |
| prompt: Detailed text description of the image to generate | |
| mode: 'fast' for quick generation (z-image turbo), 'pro' for higher quality (premium models) | |
| model: Model for pro mode - 'seedream4', 'imagen4_ultra', 'flux-krea' | |
| width: Image width in pixels | |
| height: Image height in pixels | |
| num_images: Number of images to generate (1-4) | |
| negative_prompt: Things to avoid in the generated image | |
| Returns: | |
| String containing the file paths of generated images | |
| """ | |
| # Clamp values | |
| width = max(256, min(2048, width)) | |
| height = max(256, min(2048, height)) | |
| num_images = max(1, min(4, num_images)) | |
| # Get model ID for pro mode | |
| model_id = PRO_MODELS.get(model, PRO_MODELS["flux-krea"]) | |
| # Debug logging | |
| debug_info = { | |
| "prompt": prompt[:100] + "..." if len(prompt) > 100 else prompt, | |
| "mode": mode, | |
| "model": model if mode == "pro" else "z-image-turbo", | |
| "model_id": model_id if mode == "pro" else "fal-ai/z-image/turbo", | |
| "width": width, | |
| "height": height, | |
| "num_images": num_images, | |
| } | |
| print(f"\n{'='*80}") | |
| print(f"DEBUG - Image Generation Tool Called:") | |
| print(json.dumps(debug_info, indent=2)) | |
| print(f"{'='*80}\n") | |
| st = time.time() | |
| generated_paths = [] | |
| try: | |
| if mode == "fast": | |
| generated_paths = _generate_fast(prompt, width, height, num_images) | |
| else: | |
| generated_paths = _generate_pro(prompt, model_id, width, height, num_images) | |
| except Exception as e: | |
| et = time.time() | |
| print(f"DEBUG - Image generation FAILED after {et - st:.2f}s: {e}") | |
| return f"Failed to generate images: {str(e)}" | |
| et = time.time() | |
| if not generated_paths: | |
| print(f"DEBUG - Image generation FAILED after {et - st:.2f}s: No images returned") | |
| return "Failed to generate images: No images returned from API" | |
| print(f"DEBUG - Generated {len(generated_paths)} images in {et - st:.2f}s") | |
| for path in generated_paths: | |
| print(f"DEBUG - Image path: {path}") | |
| result = f"Successfully generated {len(generated_paths)} image(s): {', '.join(generated_paths)}" | |
| print(f"DEBUG - Image generation result:\n{result}") | |
| return result | |