"""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: " 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