Spaces:
Running
Running
| import argparse | |
| import os | |
| import sys | |
| import requests | |
| from dotenv import load_dotenv | |
| from stream_utils import StreamResponseHandler | |
| # Load environment variables | |
| load_dotenv() | |
| def get_default_image(): | |
| """Get the default image path and convert it to a data URI.""" | |
| image_path = "./assets/lakeview.jpg" | |
| if os.path.exists(image_path): | |
| try: | |
| from src.utils import image_path_to_uri | |
| image_uri = image_path_to_uri(image_path) | |
| print(f"Using default image: {image_path}") | |
| return image_uri | |
| except Exception as e: | |
| print(f"Error converting image to URI: {e}") | |
| return None | |
| else: | |
| print(f"Warning: Default image not found at {image_path}") | |
| return None | |
| def upload_image(handler, image_path): | |
| """ | |
| Upload an image to the server. | |
| Args: | |
| handler (StreamResponseHandler): The stream response handler. | |
| image_path (str): Path to the image file to upload. | |
| Returns: | |
| str: The URL of the uploaded image, or None if upload failed. | |
| """ | |
| if not os.path.exists(image_path): | |
| handler.console.print( | |
| f"[bold red]Error:[/] Image file not found at {image_path}" | |
| ) | |
| return None | |
| try: | |
| handler.console.print(f"Uploading image: [bold]{image_path}[/]") | |
| with open(image_path, "rb") as f: | |
| files = {"file": (os.path.basename(image_path), f)} | |
| response = requests.post("http://localhost:8000/upload", files=files) | |
| if response.status_code == 200: | |
| image_url = response.json().get("image_url") | |
| handler.console.print( | |
| f"Image uploaded successfully. URL: [bold green]{image_url}[/]" | |
| ) | |
| return image_url | |
| else: | |
| handler.console.print( | |
| f"[bold red]Failed to upload image.[/] Status code: {response.status_code}" | |
| ) | |
| handler.console.print(f"Response: {response.text}") | |
| return None | |
| except Exception as e: | |
| handler.console.print(f"[bold red]Error uploading image:[/] {e}") | |
| return None | |
| def main(): | |
| # Create a stream response handler | |
| handler = StreamResponseHandler() | |
| # Parse command line arguments | |
| parser = argparse.ArgumentParser(description="Test the image edit streaming API.") | |
| parser.add_argument( | |
| "--instruction", "-i", required=True, help="The edit instruction." | |
| ) | |
| parser.add_argument("--image", "-img", help="The URL of the image to edit.") | |
| parser.add_argument("--upload", "-u", help="Path to an image file to upload first.") | |
| args = parser.parse_args() | |
| # Check if the server is running | |
| if not handler.check_server_health(): | |
| sys.exit(1) | |
| image_url = args.image | |
| # If upload is specified, upload the image first | |
| if args.upload: | |
| image_url = upload_image(handler, args.upload) | |
| if not image_url: | |
| handler.console.print( | |
| "[yellow]Warning:[/] Failed to upload image. Continuing without image URL." | |
| ) | |
| # Use the default image if no image URL is provided | |
| if not image_url: | |
| image_url = get_default_image() | |
| if not image_url: | |
| handler.console.print( | |
| "[yellow]No image URL provided and default image not available.[/]" | |
| ) | |
| handler.console.print("The agent may ask for an image if needed.") | |
| # Prepare the payload for the edit request | |
| payload = {"edit_instruction": args.instruction} | |
| if image_url: | |
| payload["image_url"] = image_url | |
| # Stream the edit request | |
| endpoint_url = "http://localhost:8000/edit/stream" | |
| handler.stream_response(endpoint_url, payload=payload, title="Image Edit Response") | |
| if __name__ == "__main__": | |
| main() | |