79 lines
4.2 KiB
Python
79 lines
4.2 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from http.server import BaseHTTPRequestHandler,ThreadingHTTPServer
|
||
|
|
import io,json,time
|
||
|
|
from pathlib import Path
|
||
|
|
from threading import Thread
|
||
|
|
|
||
|
|
from PIL import Image
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from spriteforge.relay.client import RelayClient
|
||
|
|
from spriteforge.relay.server import RelayServer,RelayStore
|
||
|
|
from spriteforge.relay.worker import GpuWorker
|
||
|
|
from spriteforge.studio.models import GenerationRecipe,GenerationRequest
|
||
|
|
from spriteforge.studio.providers import RelayProvider
|
||
|
|
|
||
|
|
|
||
|
|
class FakeComfy(BaseHTTPRequestHandler):
|
||
|
|
def log_message(self,*args):pass
|
||
|
|
def reply(self,value):
|
||
|
|
body=json.dumps(value).encode();self.send_response(200);self.send_header("Content-Type","application/json");self.send_header("Content-Length",str(len(body)));self.end_headers();self.wfile.write(body)
|
||
|
|
def do_POST(self):
|
||
|
|
body=self.rfile.read(int(self.headers.get("Content-Length",0)))
|
||
|
|
if self.path=="/prompt":return self.reply({"prompt_id":"p1"})
|
||
|
|
if self.path=="/upload/image":return self.reply({"name":"input.png"})
|
||
|
|
self.send_error(404)
|
||
|
|
def do_GET(self):
|
||
|
|
if self.path=="/history/p1":return self.reply({"p1":{"outputs":{"7":{"images":[{"filename":"out.png","subfolder":"","type":"output"}]}}}})
|
||
|
|
if self.path.startswith("/view?"):
|
||
|
|
image=Image.new("RGBA",(16,16),(70,20,30,255));out=io.BytesIO();image.save(out,"PNG");data=out.getvalue();self.send_response(200);self.send_header("Content-Length",str(len(data)));self.end_headers();self.wfile.write(data);return
|
||
|
|
if self.path=="/system_stats":return self.reply({"ok":True})
|
||
|
|
self.send_error(404)
|
||
|
|
|
||
|
|
|
||
|
|
def servers(tmp_path):
|
||
|
|
relay=RelayServer(("127.0.0.1",0),RelayStore(tmp_path/"relay"),"client-secret","worker-secret");Thread(target=relay.serve_forever,daemon=True).start()
|
||
|
|
comfy=ThreadingHTTPServer(("127.0.0.1",0),FakeComfy);Thread(target=comfy.serve_forever,daemon=True).start()
|
||
|
|
return relay,comfy
|
||
|
|
|
||
|
|
|
||
|
|
def workflow(path:Path):
|
||
|
|
path.write_text(json.dumps({"1":{"class_type":"Test","inputs":{"p":"{{PROMPT}}","n":"{{NEGATIVE_PROMPT}}","seed":"{{SEED}}","w":"{{WIDTH}}","h":"{{HEIGHT}}"}}}))
|
||
|
|
|
||
|
|
|
||
|
|
def test_relay_auth_blob_lease_and_completion(tmp_path):
|
||
|
|
relay,comfy=servers(tmp_path);base=f"http://127.0.0.1:{relay.server_port}";client=RelayClient(base,"client-secret");worker=RelayClient(base,"worker-secret")
|
||
|
|
try:
|
||
|
|
digest=client.upload(b"hello");job=client.submit({"input":digest});leased=worker.lease("w1",60)
|
||
|
|
assert client.ping_client()["role"]=="client" and worker.ping_worker()["role"]=="worker"
|
||
|
|
assert leased["id"]==job["id"] and leased["attempts"]==1
|
||
|
|
output=worker.upload(b"result");worker.complete(job["id"],"w1",[output])
|
||
|
|
assert client.status(job["id"])["status"]=="succeeded" and client.download(output)==b"result"
|
||
|
|
with pytest.raises(ValueError,match="401"):RelayClient(base,"wrong").status(job["id"])
|
||
|
|
queued=client.submit({"x":1});client.cancel(queued["id"]);assert client.status(queued["id"])["status"]=="cancelled"
|
||
|
|
finally:relay.shutdown();relay.server_close();comfy.shutdown();comfy.server_close()
|
||
|
|
|
||
|
|
|
||
|
|
def test_studio_relay_provider_to_gpu_worker_to_comfyui(tmp_path):
|
||
|
|
relay,comfy=servers(tmp_path);base=f"http://127.0.0.1:{relay.server_port}";flow=tmp_path/"workflow.json";workflow(flow)
|
||
|
|
provider=RelayProvider(base,flow,"client-secret",timeout=5,poll_interval=.01)
|
||
|
|
request=GenerationRequest(project_id="ops",asset_id="agent",shot_id="shot",count=1,
|
||
|
|
recipe=GenerationRecipe(provider="relay",model="test",seed=3,prompt="horror",width=64,height=64))
|
||
|
|
result=[];errors=[]
|
||
|
|
thread=Thread(target=lambda: _capture(provider,request,result,errors));thread.start()
|
||
|
|
worker=GpuWorker(base,"worker-secret",f"http://127.0.0.1:{comfy.server_port}","gpu1",60)
|
||
|
|
try:
|
||
|
|
deadline=time.monotonic()+3
|
||
|
|
while time.monotonic()<deadline:
|
||
|
|
if worker.once():break
|
||
|
|
time.sleep(.01)
|
||
|
|
thread.join(5)
|
||
|
|
assert not errors and result and result[0][0].startswith(b"\x89PNG")
|
||
|
|
finally:relay.shutdown();relay.server_close();comfy.shutdown();comfy.server_close()
|
||
|
|
|
||
|
|
|
||
|
|
def _capture(provider,request,result,errors):
|
||
|
|
try:result.append(provider.generate(request,lambda _:b""))
|
||
|
|
except Exception as error:errors.append(error)
|