DN6
66 lines
1.8 KiB
Python
66 lines
1.8 KiB
Python
import os
|
|
import time
|
|
import urllib.parse
|
|
|
|
import requests
|
|
|
|
|
|
SERVER_URL = "http://localhost:8500/api/diffusers/inference"
|
|
BASE_URL = "http://localhost:8500"
|
|
DOWNLOAD_FOLDER = "generated_images"
|
|
WAIT_BEFORE_DOWNLOAD = 2 # seconds
|
|
|
|
os.makedirs(DOWNLOAD_FOLDER, exist_ok=True)
|
|
|
|
|
|
def save_from_url(url: str) -> str:
|
|
"""Download the given URL (relative or absolute) and save it locally."""
|
|
if url.startswith("/"):
|
|
direct = BASE_URL.rstrip("/") + url
|
|
else:
|
|
direct = url
|
|
resp = requests.get(direct, timeout=60)
|
|
resp.raise_for_status()
|
|
filename = os.path.basename(urllib.parse.urlparse(direct).path) or f"img_{int(time.time())}.png"
|
|
path = os.path.join(DOWNLOAD_FOLDER, filename)
|
|
with open(path, "wb") as f:
|
|
f.write(resp.content)
|
|
return path
|
|
|
|
|
|
def main():
|
|
payload = {
|
|
"prompt": "The T-800 Terminator Robot Returning From The Future, Anime Style",
|
|
"num_inference_steps": 30,
|
|
"num_images_per_prompt": 1,
|
|
}
|
|
|
|
print("Sending request...")
|
|
try:
|
|
r = requests.post(SERVER_URL, json=payload, timeout=480)
|
|
r.raise_for_status()
|
|
except Exception as e:
|
|
print(f"Request failed: {e}")
|
|
return
|
|
|
|
body = r.json().get("response", [])
|
|
# Normalize to a list
|
|
urls = body if isinstance(body, list) else [body] if body else []
|
|
if not urls:
|
|
print("No URLs found in the response. Check the server output.")
|
|
return
|
|
|
|
print(f"Received {len(urls)} URL(s). Waiting {WAIT_BEFORE_DOWNLOAD}s before downloading...")
|
|
time.sleep(WAIT_BEFORE_DOWNLOAD)
|
|
|
|
for u in urls:
|
|
try:
|
|
path = save_from_url(u)
|
|
print(f"Image saved to: {path}")
|
|
except Exception as e:
|
|
print(f"Error downloading {u}: {e}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|