Spaces:
Running
Running
File size: 3,864 Bytes
2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 2a2c2ad fcb8f25 |
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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
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()
|