278 lines
11 KiB
Python
278 lines
11 KiB
Python
"""Minimal Draw Things gRPC client.
|
|
|
|
The wire contract (verified against drawthingsai/draw-things-community):
|
|
- ImageGenerationRequest.configuration is a FlatBuffers-encoded
|
|
GenerationConfiguration (Libraries/DataModels/Sources/config.fbs).
|
|
- start_width / start_height are in units of 64px (model metadata
|
|
reports default_scale=16, i.e. 1024px).
|
|
- The app's API server speaks TLS with a self-signed leaf
|
|
(CN=localhost, issuer "Draw Things Root CA"). Only the leaf is
|
|
served, so it is pinned directly as the trust root.
|
|
"""
|
|
|
|
import socket, ssl, io
|
|
import flatbuffers
|
|
import grpc
|
|
import imageService_pb2 as pb
|
|
import imageService_pb2_grpc as pbg
|
|
|
|
# --- GenerationConfiguration field slots, in config.fbs declaration order.
|
|
# Deprecated fields still consume a slot, so the ordering must be literal.
|
|
S = {
|
|
"id": 0, "start_width": 1, "start_height": 2, "seed": 3, "steps": 4,
|
|
"guidance_scale": 5, "strength": 6, "model": 7, "sampler": 8,
|
|
"batch_count": 9, "batch_size": 10, "hires_fix": 11,
|
|
"hires_fix_start_width": 12, "hires_fix_start_height": 13,
|
|
"hires_fix_strength": 14, "upscaler": 15, "image_guidance_scale": 16,
|
|
"seed_mode": 17, "clip_skip": 18, "controls": 19, "loras": 20,
|
|
"mask_blur": 21, "face_restoration": 22,
|
|
"_decode_with_attention": 23, "_hires_fix_decode_with_attention": 24,
|
|
"clip_weight": 25, "negative_prompt_for_image_prior": 26,
|
|
"image_prior_steps": 27, "refiner_model": 28,
|
|
"original_image_height": 29, "original_image_width": 30,
|
|
"crop_top": 31, "crop_left": 32, "target_image_height": 33,
|
|
"target_image_width": 34, "aesthetic_score": 35,
|
|
"negative_aesthetic_score": 36, "zero_negative_prompt": 37,
|
|
"refiner_start": 38, "negative_original_image_height": 39,
|
|
"negative_original_image_width": 40, "name": 41, "fps_id": 42,
|
|
"motion_bucket_id": 43, "cond_aug": 44, "start_frame_cfg": 45,
|
|
"num_frames": 46, "mask_blur_outset": 47, "sharpness": 48, "shift": 49,
|
|
"stage_2_steps": 50, "stage_2_cfg": 51, "stage_2_shift": 52,
|
|
"tiled_decoding": 53, "decoding_tile_width": 54,
|
|
"decoding_tile_height": 55, "decoding_tile_overlap": 56,
|
|
"stochastic_sampling_gamma": 57, "preserve_original_after_inpaint": 58,
|
|
"tiled_diffusion": 59, "diffusion_tile_width": 60,
|
|
"diffusion_tile_height": 61, "diffusion_tile_overlap": 62,
|
|
"upscaler_scale_factor": 63, "t5_text_encoder": 64,
|
|
"separate_clip_l": 65, "clip_l_text": 66, "separate_open_clip_g": 67,
|
|
"open_clip_g_text": 68, "speed_up_with_guidance_embed": 69,
|
|
"guidance_embed": 70, "resolution_dependent_shift": 71,
|
|
"tea_cache_start": 72, "tea_cache_end": 73, "tea_cache_threshold": 74,
|
|
"tea_cache": 75, "separate_t5": 76, "t5_text": 77,
|
|
"tea_cache_max_skip_steps": 78, "causal_inference_enabled": 79,
|
|
"causal_inference": 80, "causal_inference_pad": 81, "cfg_zero_star": 82,
|
|
"cfg_zero_init_steps": 83, "compression_artifacts": 84,
|
|
"compression_artifacts_quality": 85, "color_calibration": 86,
|
|
"expand_prompt_to_json": 87,
|
|
}
|
|
NFIELDS = 88
|
|
|
|
SAMPLER = {
|
|
"DPMPP2MKarras": 0, "EulerA": 1, "DDIM": 2, "PLMS": 3,
|
|
"DPMPPSDEKarras": 4, "UniPC": 5, "LCM": 6, "EulerASubstep": 7,
|
|
"DPMPPSDESubstep": 8, "TCD": 9, "EulerATrailing": 10,
|
|
"DPMPPSDETrailing": 11, "DPMPP2MAYS": 12, "EulerAAYS": 13,
|
|
"DPMPPSDEAYS": 14, "DPMPP2MTrailing": 15, "DDIMTrailing": 16,
|
|
"UniPCTrailing": 17, "UniPCAYS": 18, "TCDTrailing": 19,
|
|
}
|
|
SEED_MODE = {"Legacy": 0, "TorchCpuCompatible": 1, "ScaleAlike": 2,
|
|
"NvidiaGpuCompatible": 3}
|
|
|
|
|
|
def build_config(model, width, height, steps, guidance_scale, seed,
|
|
sampler="DPMPP2MTrailing", batch_count=1, shift=None,
|
|
strength=1.0, loras=(), seed_mode="ScaleAlike",
|
|
clip_skip=1, name=None):
|
|
"""Serialize a GenerationConfiguration FlatBuffer.
|
|
|
|
width/height are in PIXELS here and converted to 64px units.
|
|
loras is an iterable of (filename, weight).
|
|
"""
|
|
if width % 64 or height % 64:
|
|
raise ValueError("width/height must be multiples of 64")
|
|
|
|
b = flatbuffers.Builder(4096)
|
|
|
|
# All strings/vectors must be built before StartObject.
|
|
model_off = b.CreateString(model)
|
|
name_off = b.CreateString(name) if name else None
|
|
|
|
lora_offs = []
|
|
for f, w in loras:
|
|
f_off = b.CreateString(f)
|
|
b.StartObject(3) # LoRA: file, weight, mode
|
|
b.PrependUOffsetTRelativeSlot(0, f_off, 0)
|
|
b.PrependFloat32Slot(1, float(w), 0.6)
|
|
lora_offs.append(b.EndObject())
|
|
loras_vec = None
|
|
if lora_offs:
|
|
b.StartVector(4, len(lora_offs), 4)
|
|
for off in reversed(lora_offs):
|
|
b.PrependUOffsetTRelative(off)
|
|
loras_vec = b.EndVector()
|
|
|
|
b.StartObject(NFIELDS)
|
|
b.PrependInt64Slot(S["id"], 0, 0)
|
|
b.PrependUint16Slot(S["start_width"], width // 64, 0)
|
|
b.PrependUint16Slot(S["start_height"], height // 64, 0)
|
|
b.PrependUint32Slot(S["seed"], int(seed) & 0xFFFFFFFF, 0)
|
|
b.PrependUint32Slot(S["steps"], int(steps), 0)
|
|
b.PrependFloat32Slot(S["guidance_scale"], float(guidance_scale), 0.0)
|
|
b.PrependFloat32Slot(S["strength"], float(strength), 0.0)
|
|
b.PrependUOffsetTRelativeSlot(S["model"], model_off, 0)
|
|
b.PrependInt8Slot(S["sampler"], SAMPLER[sampler], 0)
|
|
b.PrependUint32Slot(S["batch_count"], int(batch_count), 1)
|
|
b.PrependUint32Slot(S["batch_size"], 1, 1)
|
|
b.PrependBoolSlot(S["hires_fix"], False, False)
|
|
b.PrependInt8Slot(S["seed_mode"], SEED_MODE[seed_mode], 0)
|
|
b.PrependUint32Slot(S["clip_skip"], int(clip_skip), 1)
|
|
if loras_vec is not None:
|
|
b.PrependUOffsetTRelativeSlot(S["loras"], loras_vec, 0)
|
|
b.PrependFloat32Slot(S["clip_weight"], 1.0, 1.0)
|
|
if shift is not None:
|
|
b.PrependFloat32Slot(S["shift"], float(shift), 1.0)
|
|
b.PrependBoolSlot(S["resolution_dependent_shift"], False, True)
|
|
if name_off is not None:
|
|
b.PrependUOffsetTRelativeSlot(S["name"], name_off, 0)
|
|
cfg = b.EndObject()
|
|
b.Finish(cfg)
|
|
return bytes(b.Output())
|
|
|
|
|
|
def connect(host="127.0.0.1", port=7859, timeout=10):
|
|
"""Open a channel, pinning the server's self-signed leaf certificate."""
|
|
ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
|
|
ctx.check_hostname = False
|
|
ctx.verify_mode = ssl.CERT_NONE
|
|
ctx.set_alpn_protocols(["h2"])
|
|
with socket.create_connection((host, port), timeout=timeout) as raw:
|
|
with ctx.wrap_socket(raw, server_hostname="localhost") as s:
|
|
pem = ssl.DER_cert_to_PEM_cert(s.getpeercert(binary_form=True))
|
|
|
|
creds = grpc.ssl_channel_credentials(root_certificates=pem.encode())
|
|
ch = grpc.secure_channel(f"{host}:{port}", creds, options=[
|
|
("grpc.ssl_target_name_override", "localhost"),
|
|
("grpc.max_receive_message_length", 512 * 1024 * 1024),
|
|
("grpc.max_send_message_length", 512 * 1024 * 1024),
|
|
])
|
|
return pbg.ImageGenerationServiceStub(ch)
|
|
|
|
|
|
CCV_MAGIC = 0x000F7217
|
|
CCV_16F, CCV_32F, CCV_8U = 0x20000, 0x40000, 0x1000
|
|
|
|
|
|
def decode_tensor(blob):
|
|
"""Decode an s4nnc tensor blob into a uint8 HxWxC numpy array.
|
|
|
|
Layout (verified empirically against the running server):
|
|
u32 magic (0x000f7217) | u32 memory type | u32 format (2=NHWC)
|
|
u32 datatype | u32 reserved | 12 x u32 dims -> 68-byte header
|
|
followed by the payload, which is fpzip-compressed (magic "fpy"),
|
|
zlib-compressed, or raw depending on the server's codec choice.
|
|
Pixel values are float in [-1, 1].
|
|
"""
|
|
import struct
|
|
import numpy as np
|
|
|
|
magic, mtype, fmt, dtype, _ = struct.unpack_from("<5I", blob, 0)
|
|
if magic != CCV_MAGIC:
|
|
raise ValueError("not a ccv tensor blob: magic=0x%08x" % magic)
|
|
dims = [d for d in struct.unpack_from("<12I", blob, 20) if d]
|
|
payload = blob[68:]
|
|
|
|
if payload[:3] == b"fpy":
|
|
import fpzip
|
|
arr = fpzip.decompress(payload).astype("float32")
|
|
else:
|
|
if payload[:1] == b"\x78":
|
|
import zlib
|
|
payload = zlib.decompress(payload)
|
|
np_dtype = {CCV_16F: "<f2", CCV_32F: "<f4", CCV_8U: "u1"}[dtype]
|
|
arr = np.frombuffer(payload, dtype=np_dtype).astype("float32")
|
|
|
|
arr = arr.reshape(dims)
|
|
while arr.ndim > 3 and arr.shape[0] == 1: # drop leading batch dim
|
|
arr = arr[0]
|
|
if dtype == CCV_8U:
|
|
return arr.astype("uint8")
|
|
return np.clip((arr + 1.0) * 127.5, 0, 255).astype("uint8")
|
|
|
|
|
|
def decode_png(blob):
|
|
"""Decode a generated blob and return PNG bytes."""
|
|
from PIL import Image
|
|
arr = decode_tensor(blob)
|
|
buf = io.BytesIO()
|
|
Image.fromarray(arr).save(buf, format="PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
def generate(stub, prompt, negative_prompt="", config=None, timeout=3600,
|
|
on_progress=None, image=None, mask=None):
|
|
"""Run GenerateImage and return a list of image byte blobs.
|
|
|
|
image/mask are optional uint8 numpy arrays (or pre-encoded tensor blobs).
|
|
They are content-addressed: the request carries the 32-byte SHA-256 digest
|
|
in .image/.mask and the blob itself in .contents, which is how the server
|
|
resolves them (see ImageGenerationServiceImpl.unwrapData).
|
|
"""
|
|
contents = []
|
|
|
|
def attach(value):
|
|
if value is None:
|
|
return None
|
|
blob = value if isinstance(value, (bytes, bytearray)) \
|
|
else encode_tensor(value)
|
|
digest, content = content_request(bytes(blob))
|
|
contents.append(content)
|
|
return digest
|
|
|
|
image_digest = attach(image)
|
|
mask_digest = attach(mask)
|
|
|
|
req = pb.ImageGenerationRequest(
|
|
prompt=prompt,
|
|
negativePrompt=negative_prompt,
|
|
configuration=config,
|
|
scaleFactor=1,
|
|
user="sideCards",
|
|
device=pb.LAPTOP,
|
|
chunked=False,
|
|
contents=contents,
|
|
)
|
|
if image_digest is not None:
|
|
req.image = image_digest
|
|
if mask_digest is not None:
|
|
req.mask = mask_digest
|
|
|
|
images, chunks = [], []
|
|
for resp in stub.GenerateImage(req, timeout=timeout):
|
|
if on_progress and resp.HasField("currentSignpost"):
|
|
on_progress(resp.currentSignpost)
|
|
for img in resp.generatedImages:
|
|
images.append(img)
|
|
if resp.chunkState == pb.MORE_CHUNKS:
|
|
chunks.append(resp)
|
|
return images
|
|
|
|
|
|
def encode_tensor(arr):
|
|
"""Encode a uint8 HxWxC image into an s4nnc tensor blob (fpzip payload).
|
|
|
|
Mirrors decode_tensor: 68-byte ccv header then an fpzip stream, with
|
|
pixel values scaled into [-1, 1]. Datatype is marked CCV_16F to match
|
|
what the server itself emits.
|
|
"""
|
|
import struct
|
|
import numpy as np
|
|
import fpzip
|
|
|
|
a = np.asarray(arr)
|
|
if a.dtype == np.uint8:
|
|
a = a.astype("float32") / 127.5 - 1.0
|
|
a = a.astype("float32")
|
|
if a.ndim == 3:
|
|
a = a[None, ...] # NHWC with N=1
|
|
dims = list(a.shape) + [0] * (12 - a.ndim)
|
|
|
|
header = struct.pack("<5I", CCV_MAGIC, 1, 2, CCV_16F, 0)
|
|
header += struct.pack("<12I", *dims)
|
|
return header + fpzip.compress(np.ascontiguousarray(a), precision=0)
|
|
|
|
|
|
def content_request(blob):
|
|
"""Return (sha256_digest, blob) for content-addressed request fields."""
|
|
import hashlib
|
|
return hashlib.sha256(blob).digest(), blob
|