Spaces:
Running on Zero
Running on Zero
File size: 2,990 Bytes
b701455 | 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 |
import subprocess
import time
import requests
import os
import base64
import signal
import sys
def test_server():
print("Starting server...")
# Start server in a new process group to avoid being killed by window-CLOSE
server_proc = subprocess.Popen(
[sys.executable, "server.py"],
stdout=open("server_test_stdout.log", "w"),
stderr=open("server_test_stderr.log", "w"),
creationflags=subprocess.CREATE_NEW_PROCESS_GROUP
)
try:
# Wait for server to be ready
print("Waiting for server to be ready...")
ready = False
for i in range(30):
try:
resp = requests.get("http://localhost:7861/api/telemetry", timeout=2)
if resp.status_code == 200:
ready = True
print("Server is ready!")
break
except:
time.sleep(2)
if not ready:
print("Server failed to start in time.")
return
# Send generation request
print("Sending generation request...")
payload = {
"prompt": "a high quality portrait of a cat",
"width": 512,
"height": 512,
"steps": 4,
"cfg_scale": 1.0,
"sampler": "euler",
"scheduler": "ays",
"model_path": "./include/diffusion_model/flux-2-klein-4b.safetensors"
}
try:
# High timeout for generation
resp = requests.post("http://localhost:7861/api/generate", json=payload, timeout=600)
if resp.status_code == 200:
data = resp.json()
if "image" in data:
img_str = data["image"]
if "," in img_str:
img_str = img_str.split(",")[1]
try:
with open("test_flux_result.png", "wb") as f:
f.write(base64.b64decode(img_str))
print("Success! Image saved to test_flux_result.png")
except Exception as b64e:
print(f"Base64 decode failed: {b64e}")
print(f"Image string prefix: {img_str[:100]}...")
print(f"Image string suffix: {img_str[-100:]}...")
print(f"Image string length: {len(img_str)}")
else:
print(f"Error: No image in response: {data}")
else:
print(f"Error: Server returned status code {resp.status_code}")
print(resp.text)
except Exception as e:
print(f"Request failed: {e}")
finally:
print("Shutting down server...")
# Send CTRL_BREAK_EVENT to the process group
os.kill(server_proc.pid, signal.CTRL_BREAK_EVENT)
server_proc.wait(timeout=10)
if __name__ == "__main__":
test_server()
|