File size: 6,243 Bytes
1e103b7 | 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 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | import argparse
import sys
import os
import base64
import time
import io
import requests
from PIL import Image
# Oh, hello there! Nikola here, ready to help this little client talk to the big server!
# It's like sending a messenger bird from our village to the capital!
def main():
# Peeking at the arguments... gotta make sure we have all our supplies for the journey!
parser = argparse.ArgumentParser(description="ImageGen Client - A little seeker tool!")
parser.add_argument("--host", type=str, default="localhost", help="Where the server lives (Host)")
parser.add_argument("--port", type=int, default=8000, help="The door to knock on (Port)")
parser.add_argument("--num_images", type=int, default=1, help="How many pictures to paint?")
parser.add_argument("--image_folder", type=str, default="generated_images", help="Where to keep our treasures")
# Changing defaults to None so we can use input image size if needed!
parser.add_argument("--width", type=int, default=None, help="Canvas width (default: 1024 or input image size)")
parser.add_argument("--height", type=int, default=None, help="Canvas height (default: 1024 or input image size)")
# New shiny tools for our quest!
parser.add_argument("--input", type=str, default=None, help="Path to an input image (for image-to-image magic!)")
parser.add_argument("--max-size", type=int, default=1024, help="Max size for the input image (we don't want it to get too heavy for the bird!)")
args = parser.parse_args()
# Reading the prompt from the spirits... I mean, stdin!
# "What do you desire to see?" *sparkle*
print("Waiting for a prompt from stdin... (Type something and press Ctrl+D!)")
try:
prompt = sys.stdin.read().strip()
except Exception as e:
print(f"Oh no! The spirits were silent (stdin error): {e}")
return
if not prompt:
print("Aww, the prompt was empty! The canvas remains blank.")
return
print(f"Yay! We got a prompt: '{prompt}'")
# Restoring the canvas size variables from the journey's start!
final_width = args.width
final_height = args.height
# Prepare prompt and payload
url_gen = f"http://{args.host}:{args.port}/v1/images/generations"
url_edit = f"http://{args.host}:{args.port}/v1/images/edits"
try:
if args.input:
print(f"Oh! You brought a reference image: {args.input}. Let's go to the Editing Shrine!")
# Prepare for multipart upload
# We need to open the image file effectively
if not os.path.exists(args.input):
print(f"Eek! I can't find the image at {args.input}")
return
# Open image to ensure it's valid and memory-friendly resize if needed
with Image.open(args.input) as img:
img = img.convert("RGB")
w, h = img.size
max_dim = max(w, h)
if max_dim > args.max_size:
scale = args.max_size / max_dim
new_w = int(w * scale)
new_h = int(h * scale)
print(f"Resizing big image from {w}x{h} to {new_w}x{new_h}. Compact and cute!")
img = img.resize((new_w, new_h), Image.LANCZOS)
if final_width is None: final_width = img.width
if final_height is None: final_height = img.height
# Save to buffer for upload
buffered = io.BytesIO()
img.save(buffered, format="PNG")
buffered.seek(0)
image_bytes = buffered.getvalue()
# Construct multipart payload
files = {
'image': ('input.png', image_bytes, 'image/png')
}
data = {
'prompt': prompt,
'n': args.num_images,
'size': f"{final_width}x{final_height}",
'response_format': 'b64_json',
'guidance_scale': 2.5 # Default specific to edit/kontext if needed
}
print(f"Sending input image to {url_edit}... *whoosh*")
response = requests.post(url_edit, files=files, data=data)
else:
# Standard Generation
print("Just a prompt? Off to the Creation Forge!")
if final_width is None: final_width = 1024
if final_height is None: final_height = 1024
payload = {
"prompt": prompt,
"n": args.num_images,
"size": f"{final_width}x{final_height}",
"response_format": "b64_json"
}
print(f"Sending prompt to {url_gen}... *sparkle*")
response = requests.post(url_gen, json=payload)
response.raise_for_status()
data = response.json()
# Making sure we have a chest for our treasures
if not os.path.exists(args.image_folder):
print(f"Creating a new treasure chest at {args.image_folder}...")
os.makedirs(args.image_folder)
# Unpacking the magic
images = data.get("data", [])
print(f"Ooh! The server sent back {len(images)} masterpieces!")
for i, img_data in enumerate(images):
# Decoding the spell
img_bytes = base64.b64decode(img_data["b64_json"])
timestamp = int(time.time())
filename = f"image_{timestamp}_{i}.png"
filepath = os.path.join(args.image_folder, filename)
with open(filepath, "wb") as f:
f.write(img_bytes)
print(f"Saved masterpiece #{i+1} to {filepath}! It sparkles!")
except requests.exceptions.ConnectionError:
print("Oh no! The server didn't answer. Is it sleeping? (Connection Refused)")
print("Maybe check if the host and port are correct? We tried: " + url)
except Exception as e:
print(f"Eek! Something went wrong on the journey: {e}")
# We'll give it a gentle hug and try to understand...
print("Don't worry, we can try again later!")
if __name__ == "__main__":
main()
|