Spaces:
Running
Running
File size: 3,764 Bytes
2a2c2ad |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
import argparse
import os
import sys
import requests
import json
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() |