drawthings-py 0.1.0__py3-none-any.whl

drawthings_py/__init__.py

@@ -0,0 +1,19 @@
+"""
+Draw Things Python SDK
+"""
+
+import drawthings_py.drawthings as DrawThings
+from .configs import Configs, ConfigDict, Presets
+from .filename_pattern import FilenamePattern
+from .image_buffer import ImageBuffer
+from .request_builder import RequestBuilder
+
+__all__ = [
+    "DrawThings",
+    "RequestBuilder",
+    "ImageBuffer",
+    "ConfigDict",
+    "Configs",
+    "Presets",
+    "FilenamePattern",
+]

drawthings_py/_dt_service.py

@@ -0,0 +1,39 @@
+"""
+Primary entry point for using Draw Things services
+"""
+
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from types import TracebackType
+
+from .image_buffer import ImageBuffer
+from .request_builder import RequestBuilder
+
+
+class DrawThingsService(ABC):
+    """
+    Base class for grpc and cii service
+    """
+
+    @abstractmethod
+    async def generate_image(self, request: RequestBuilder) -> list[ImageBuffer]:
+        """
+        Generate an image from the provided request builder
+        """
+
+    @abstractmethod
+    def _dispose(self):
+        """
+        dispose of the service
+        """
+
+    async def __aenter__(self) -> "DrawThingsService":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        self._dispose()

drawthings_py/_errors.py

@@ -0,0 +1,21 @@
+from grpclib import GRPCError, Status
+
+
+class DrawThingsServerError(Exception):
+    pass
+
+
+class DrawThingsUnavailableError(Exception):
+    pass
+
+
+def raise_grpc_error(e: GRPCError):
+    if e.status == Status.INTERNAL:
+        raise DrawThingsServerError(
+            f"There was an error on the server. This could be a temporary issue with DT+ or a problem with your request. ({e.status.name}: {e.message})"
+        ) from e
+
+    if e.status == Status.UNAVAILABLE:
+        raise DrawThingsUnavailableError(
+            "The gRPC server is currently unavailable."
+        ) from e

drawthings_py/_metadata.py

@@ -0,0 +1,338 @@
+from __future__ import annotations
+
+import copy
+import sys
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, Required, TypedDict, cast
+
+from .generated.dt_grpc.config_generated import GenerationConfiguration
+
+SAMPLER_NAMES: dict[Any, str] = {
+    0: "DPM++ 2M Karras",
+    1: "Euler A",
+    2: "DDIM",
+    3: "PLMS",
+    4: "DPM++ SDE Karras",
+    5: "UniPC",
+    6: "LCM",
+    7: "Euler A Substep",
+    8: "DPM++ SDE Substep",
+    9: "TCD",
+    10: "Euler A Trailing",
+    11: "DPM++ SDE Trailing",
+    12: "DPM++ 2M AYS",
+    13: "Euler A AYS",
+    14: "DPM++ SDE AYS",
+    15: "DPM++ 2M Trailing",
+    16: "DDIM Trailing",
+    17: "UniPC Trailing",
+    18: "UniPC AYS",
+    19: "TCD Trailing",
+}
+
+SEED_MODE_NAMES = {
+    0: "Legacy",
+    1: "Torch CPU Compatible",
+    2: "Scale Alike",
+    3: "Nvidia GPU Compatible",
+}
+
+COMPRESSION_ARTIFACT_NAMES = {
+    0: "disabled",
+    1: "h264",
+    2: "h265",
+    3: "jpeg",
+}
+
+LORA_MODE_NAMES = {
+    0: "all",
+    1: "base",
+    2: "refiner",
+}
+
+CONTROL_MODE_NAMES = {
+    0: "balanced",
+    1: "prompt",
+    2: "control",
+}
+
+CONTROL_INPUT_TYPE_NAMES = {
+    0: "unspecified",
+    1: "custom",
+    2: "depth",
+    3: "canny",
+    4: "scribble",
+    5: "pose",
+    6: "normalbae",
+    7: "color",
+    8: "lineart",
+    9: "softedge",
+    10: "seg",
+    11: "inpaint",
+    12: "ip2p",
+    13: "shuffle",
+    14: "mlsd",
+    15: "tile",
+    16: "blur",
+    17: "lowquality",
+    18: "gray",
+}
+
+
+class V2(TypedDict, total=False):
+    aestheticScore: float
+    batchCount: int
+    batchSize: int
+    causalInference: bool
+    causalInferencePad: int
+    cfgZeroInitSteps: int
+    cfgZeroStar: float
+    clipLText: float
+    clipSkip: int
+    clipWeight: float
+    compressionArtifacts: str
+    compressionArtifactsQuality: int
+    controls: list[dict[str, Any]]
+    cropLeft: int
+    cropTop: int
+    decodingTileHeight: int
+    decodingTileOverlap: int
+    decodingTileWidth: int
+    diffusionTileHeight: int
+    diffusionTileOverlap: int
+    diffusionTileWidth: int
+    fps: int
+    guidanceEmbed: bool
+    guidanceScale: float
+    height: int
+    hiresFix: bool
+    hiresFixHeight: int
+    hiresFixStrength: float
+    hiresFixWidth: int
+    id: int
+    imageGuidanceScale: float
+    imagePriorSteps: int
+    loras: list[dict[str, Any]]
+    maskBlur: float
+    maskBlurOutset: int
+    motionScale: int
+    negativeAestheticScore: float
+    seed: int
+    steps: int
+    width: int
+
+
+class ImageMetadata(TypedDict, total=False):
+    c: str
+    uc: str
+    steps: Required[int]
+    sampler: Required[str]
+    scale: Required[float]
+    seed: Required[int]
+    size: Required[str]
+    model: Required[str]
+    strength: Required[float]
+    seed_mode: Required[str]
+    shift: Required[float]
+    v2: Required[V2]
+
+
+def create_metadata(
+    config: GenerationConfiguration, prompt: str, negative_prompt: str
+) -> ImageMetadata:
+    """Create Draw Things PNG metadata from the generation configuration."""
+    cfg = config
+    width = _pixels(cfg.StartWidth())
+    height = _pixels(cfg.StartHeight())
+    model = _decode_string(cfg.Model()) or ""
+    sampler = cfg.Sampler()
+    seed_mode = cfg.SeedMode()
+    v2 = _create_v2_metadata(cfg, width, height, model, sampler, seed_mode)
+
+    return ImageMetadata(
+        c=prompt or "",
+        uc=negative_prompt or "",
+        model=model,
+        sampler=SAMPLER_NAMES.get(sampler, str(sampler)),
+        scale=cfg.GuidanceScale(),
+        seed=cfg.Seed(),
+        seed_mode=SEED_MODE_NAMES.get(seed_mode, str(seed_mode)),
+        shift=cfg.Shift(),
+        size=f"{width}x{height}",
+        steps=cfg.Steps(),
+        strength=cfg.Strength(),
+        v2=v2,
+    )
+
+
+def _create_v2_metadata(
+    cfg: GenerationConfiguration,
+    width: int,
+    height: int,
+    model: str,
+    sampler: int,
+    seed_mode: int,
+) -> V2:
+    v2: dict[str, Any] = {
+        "aestheticScore": cfg.AestheticScore(),
+        "batchCount": cfg.BatchCount(),
+        "batchSize": cfg.BatchSize(),
+        "causalInference": cfg.CausalInference(),
+        "causalInferencePad": cfg.CausalInferencePad(),
+        "cfgZeroInitSteps": cfg.CfgZeroInitSteps(),
+        "cfgZeroStar": cfg.CfgZeroStar(),
+        "clipLText": cfg.ClipLText(),
+        "clipSkip": cfg.ClipSkip(),
+        "clipWeight": cfg.ClipWeight(),
+        "compressionArtifacts": COMPRESSION_ARTIFACT_NAMES.get(
+            cfg.CompressionArtifacts(),
+        ),
+        "compressionArtifactsQuality": cfg.CompressionArtifactsQuality(),
+        "controls": _controls(cfg),
+        "cropLeft": cfg.CropLeft(),
+        "cropTop": cfg.CropTop(),
+        "decodingTileHeight": _pixels(cfg.DecodingTileHeight()),
+        "decodingTileOverlap": _pixels(cfg.DecodingTileOverlap()),
+        "decodingTileWidth": _pixels(cfg.DecodingTileWidth()),
+        "diffusionTileHeight": _pixels(cfg.DiffusionTileHeight()),
+        "diffusionTileOverlap": _pixels(cfg.DiffusionTileOverlap()),
+        "diffusionTileWidth": _pixels(cfg.DiffusionTileWidth()),
+        "fps": cfg.FpsId(),
+        "guidanceEmbed": cfg.GuidanceEmbed(),
+        "guidanceScale": cfg.GuidanceScale(),
+        "guidingFrameNoise": cfg.CondAug(),
+        "height": height,
+        "hiresFix": cfg.HiresFix(),
+        "hiresFixHeight": _pixels(cfg.HiresFixStartHeight()),
+        "hiresFixStrength": cfg.HiresFixStrength(),
+        "hiresFixWidth": _pixels(cfg.HiresFixStartWidth()),
+        "id": cfg.Id(),
+        "imageGuidanceScale": cfg.ImageGuidanceScale(),
+        "imagePriorSteps": cfg.ImagePriorSteps(),
+        "loras": _loras(cfg),
+        "maskBlur": cfg.MaskBlur(),
+        "maskBlurOutset": cfg.MaskBlurOutset(),
+        "model": model,
+        "motionScale": cfg.MotionBucketId(),
+        "negativeAestheticScore": cfg.NegativeAestheticScore(),
+        "negativeOriginalImageHeight": _pixels(cfg.NegativeOriginalImageHeight()),
+        "negativeOriginalImageWidth": _pixels(cfg.NegativeOriginalImageWidth()),
+        "negativePromptForImagePrior": cfg.NegativePromptForImagePrior(),
+        "numFrames": cfg.NumFrames(),
+        "originalImageHeight": _pixels(cfg.OriginalImageHeight()),
+        "originalImageWidth": _pixels(cfg.OriginalImageWidth()),
+        "preserveOriginalAfterInpaint": cfg.PreserveOriginalAfterInpaint(),
+        "refinerStart": cfg.RefinerStart(),
+        "resolutionDependentShift": cfg.ResolutionDependentShift(),
+        "sampler": sampler,
+        "seed": cfg.Seed(),
+        "seedMode": seed_mode,
+        "separateClipL": cfg.SeparateClipL(),
+        "separateOpenClipG": cfg.SeparateOpenClipG(),
+        "separateT5": cfg.SeparateT5(),
+        "sharpness": cfg.Sharpness(),
+        "shift": _metadata_float(cfg.Shift()),
+        "speedUpWithGuidanceEmbed": cfg.SpeedUpWithGuidanceEmbed(),
+        "stage2Guidance": cfg.Stage2Cfg(),
+        "stage2Shift": cfg.Stage2Shift(),
+        "stage2Steps": cfg.Stage2Steps(),
+        "startFrameGuidance": cfg.StartFrameCfg(),
+        "steps": cfg.Steps(),
+        "stochasticSamplingGamma": _metadata_float(cfg.StochasticSamplingGamma()),
+        "strength": cfg.Strength(),
+        "t5TextEncoder": cfg.T5TextEncoder(),
+        "targetImageHeight": _pixels(cfg.TargetImageHeight()),
+        "targetImageWidth": _pixels(cfg.TargetImageWidth()),
+        "teaCache": cfg.TeaCache(),
+        "teaCacheEnd": cfg.TeaCacheEnd(),
+        "teaCacheMaxSkipSteps": cfg.TeaCacheMaxSkipSteps(),
+        "teaCacheStart": cfg.TeaCacheStart(),
+        "teaCacheThreshold": cfg.TeaCacheThreshold(),
+        "tiledDecoding": cfg.TiledDecoding(),
+        "tiledDiffusion": cfg.TiledDiffusion(),
+        "upscalerScaleFactor": cfg.UpscalerScaleFactor(),
+        "width": width,
+        "zeroNegativePrompt": cfg.ZeroNegativePrompt(),
+    }
+    v2 = {k: v for k, v in v2.items() if v is not None and v != ""}
+    return cast(V2, v2)  # pyright: ignore[reportInvalidCast]
+
+
+def _decode_string(value: bytes | str | None) -> str | None:
+    if value is None:
+        return None
+    if isinstance(value, bytes):
+        return value.decode("utf-8")
+    return value
+
+
+def _pixels(value: int | None) -> int:
+    return (value or 0) * 64
+
+
+def _metadata_float(value: float) -> float:
+    return round(float(value), 7)
+
+
+def _loras(cfg: GenerationConfiguration) -> list[dict[str, Any]]:
+    return [
+        {
+            "mode": LORA_MODE_NAMES.get(lora.Mode(), str(lora.Mode())),
+            "file": _decode_string(lora.File()) or "",
+            "weight": lora.Weight(),
+        }
+        for lora in _nested_items(cfg, "Loras", "LorasLength")
+    ]
+
+
+def _controls(cfg: GenerationConfiguration) -> list[dict[str, Any]]:
+    controls = []
+    for control in _nested_items(cfg, "Controls", "ControlsLength"):
+        controls.append(
+            {
+                "file": _decode_string(control.File()) or "",
+                "weight": control.Weight(),
+                "guidanceStart": control.GuidanceStart(),
+                "guidanceEnd": control.GuidanceEnd(),
+                "noPrompt": control.NoPrompt(),
+                "globalAveragePooling": control.GlobalAveragePooling(),
+                "downSamplingRate": control.DownSamplingRate(),
+                "controlMode": CONTROL_MODE_NAMES.get(
+                    control.ControlMode(), str(control.ControlMode())
+                ),
+                "targetBlocks": [
+                    _decode_string(control.TargetBlocks(index)) or ""
+                    for index in range(control.TargetBlocksLength())
+                ],
+                "inputOverride": CONTROL_INPUT_TYPE_NAMES.get(
+                    control.InputOverride(), str(control.InputOverride())
+                ),
+            }
+        )
+    return controls
+
+
+def _nested_items(
+    cfg: GenerationConfiguration,
+    item_name: str,
+    length_name: str,
+) -> list[Any]:
+    generated_dir = str(Path(__file__).parent / "generated" / "dt_grpc")
+    added_path = generated_dir not in sys.path
+    if added_path:
+        sys.path.append(generated_dir)
+    try:
+        item: Callable[[int], Any] = getattr(cfg, item_name)
+        length: Callable[[], int] = getattr(cfg, length_name)
+        return [item(index) for index in range(length())]
+    finally:
+        if added_path:
+            sys.path.remove(generated_dir)
+
+
+def copy_with_seed(metadata: ImageMetadata, seed: int) -> ImageMetadata:
+    new_metadata = copy.deepcopy(metadata)
+    new_metadata["seed"] = seed
+    new_metadata["v2"]["seed"] = seed
+    return new_metadata

drawthings_py/_png_writer.py

@@ -0,0 +1,262 @@
+"""
+Utils for writing Draw Things metadata to pngs
+"""
+
+import io
+import json
+import struct
+import zlib
+from typing import cast
+from typing_extensions import override
+
+from PIL import Image
+
+from ._metadata import ImageMetadata
+
+
+class BytesEncoder(json.JSONEncoder):
+    """Custom JSON encoder that handles bytes by decoding them to UTF-8 strings."""
+
+    @override
+    def default(self, o: object) -> str:  # noqa: ANN401
+        if isinstance(o, bytes):
+            try:
+                return o.decode("utf-8")
+            except UnicodeDecodeError:
+                return o.hex()
+        return cast(str, super().default(o))
+
+
+PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
+
+
+def png_chunk(chunk_type: bytes, data: bytes) -> bytes:
+    crc = zlib.crc32(chunk_type)
+    crc = zlib.crc32(data, crc)
+
+    return (
+        struct.pack(">I", len(data))
+        + chunk_type
+        + data
+        + struct.pack(">I", crc & 0xFFFFFFFF)
+    )
+
+
+def build_itxt_chunk(keyword: str, text: str) -> bytes:
+    out = bytearray()
+
+    out.extend(keyword.encode("utf-8"))
+    out.append(0)  # keyword null terminator
+
+    out.append(0)  # compression flag
+    out.append(0)  # compression method
+
+    out.append(0)  # language tag empty
+    out.append(0)  # translated keyword empty
+
+    out.extend(text.encode("utf-8"))
+
+    return bytes(out)
+
+
+def build_exif_user_comment(width: int, height: int) -> bytes:
+    exif = bytearray()
+
+    # TIFF header
+    exif.extend(b"MM")  # big endian
+    exif.extend(struct.pack(">H", 42))
+    exif.extend(struct.pack(">I", 8))
+
+    # IFD0
+    exif.extend(struct.pack(">H", 1))
+
+    # ExifOffset tag
+    exif.extend(struct.pack(">H", 0x8769))
+    exif.extend(struct.pack(">H", 4))  # LONG
+    exif.extend(struct.pack(">I", 1))
+    exif.extend(struct.pack(">I", 26))
+
+    # next IFD
+    exif.extend(struct.pack(">I", 0))
+
+    # Exif SubIFD
+    exif.extend(struct.pack(">H", 2))
+
+    # ExifImageWidth
+    exif.extend(struct.pack(">H", 0xA002))
+    exif.extend(struct.pack(">H", 4))
+    exif.extend(struct.pack(">I", 1))
+    exif.extend(struct.pack(">I", width))
+
+    # ExifImageHeight
+    exif.extend(struct.pack(">H", 0xA003))
+    exif.extend(struct.pack(">H", 4))
+    exif.extend(struct.pack(">I", 1))
+    exif.extend(struct.pack(">I", height))
+
+    # next SubIFD
+    exif.extend(struct.pack(">I", 0))
+
+    return bytes(exif)
+
+
+def format_desc_float(f: float) -> str:
+    s = str(float(f))
+
+    if "." not in s:
+        return s + ".0"
+
+    return s
+
+
+def build_description(metadata: ImageMetadata | None) -> str:
+    if metadata is None:
+        return ""
+    desc: list[str] = []
+    if c := metadata.get("c"):
+        desc.append("Prompt: " + c)
+    if uc := metadata.get("uc"):
+        desc.append("Negative Prompt: " + uc)
+    if steps := metadata.get("steps"):
+        desc.append("Steps: " + str(steps))
+    if sampler := metadata.get("sampler"):
+        desc.append("Sampler: " + sampler)
+    if guidance_scale := metadata["v2"].get("guidanceScale", 0.0):
+        desc.append("Guidance Scale: " + format_desc_float(guidance_scale))
+    if seed := metadata.get("seed"):
+        desc.append("Seed: " + str(seed))
+    if size := metadata.get("size"):
+        desc.append("Size: " + size)
+    if model := metadata.get("model"):
+        desc.append("Model: " + model)
+    if strength := metadata.get("strength"):
+        desc.append("Strength: " + format_desc_float(strength))
+
+    return "\n".join(desc)
+
+
+def build_drawthings_xmp(json_string: str, description: str) -> str:
+    escaped_description = description.replace("\n", "
")
+
+    return f"""<x:xmpmeta xmlns:x="adobe:ns:meta/" x:xmptk="XMP Core 6.0.0">
+    <rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">
+        <rdf:Description rdf:about=""
+            xmlns:dc="http://purl.org/dc/elements/1.1/"
+            xmlns:xmp="http://ns.adobe.com/xap/1.0/"
+            xmlns:exif="http://ns.adobe.com/exif/1.0/">
+            <dc:description>
+                <rdf:Alt>
+                    <rdf:li xml:lang="x-default">{escaped_description}</rdf:li>
+                </rdf:Alt>
+            </dc:description>
+            <xmp:CreatorTool>Draw Things</xmp:CreatorTool>
+            <exif:UserComment>
+                <rdf:Alt>
+                    <rdf:li xml:lang="x-default">{json_string}</rdf:li>
+                </rdf:Alt>
+            </exif:UserComment>
+      </rdf:Description>
+   </rdf:RDF>
+</x:xmpmeta>
+"""
+
+
+def write_png_with_usercomment(
+    pixels: bytes,
+    width: int,
+    height: int,
+    channels: int,
+    metadata: ImageMetadata | None = None,
+) -> bytes:
+    mode_map = {
+        1: "L",
+        2: "LA",
+        3: "RGB",
+        4: "RGBA",
+    }
+
+    if channels not in mode_map:
+        raise ValueError("Unsupported channel count")
+
+    mode = mode_map[channels]
+
+    img = Image.frombytes(mode, (width, height), pixels)
+
+    temp = io.BytesIO()
+    img.save(temp, format="PNG")
+
+    png_data = temp.getvalue()
+
+    # Parse existing PNG chunks
+    pos = 8
+    chunks: list[tuple[bytes, bytes, bytes]] = []
+
+    while pos < len(png_data):
+        length = cast(int, struct.unpack(">I", png_data[pos : pos + 4])[0])
+        chunk_type = png_data[pos + 4 : pos + 8]
+        chunk_data = png_data[pos + 8 : pos + 8 + length]
+        chunk_crc = png_data[pos + 8 + length : pos + 12 + length]
+
+        chunks.append(
+            (
+                chunk_type,
+                chunk_data,
+                chunk_crc,
+            )
+        )
+
+        pos += length + 12
+
+    out = bytearray()
+    out.extend(PNG_SIGNATURE)
+
+    inserted = False
+
+    for chunk_type, chunk_data, chunk_crc in chunks:
+        out.extend(
+            struct.pack(">I", len(chunk_data)) + chunk_type + chunk_data + chunk_crc
+        )
+
+        if chunk_type == b"IHDR" and not inserted:
+            # sRGB chunk
+            out.extend(
+                png_chunk(
+                    b"sRGB",
+                    b"\x00",
+                )
+            )
+
+            if metadata is not None:
+                json_string = json.dumps(
+                    metadata, separators=(",", ":"), cls=BytesEncoder
+                )
+
+                exif = build_exif_user_comment(width, height)
+
+                out.extend(
+                    png_chunk(
+                        b"eXIf",
+                        exif,
+                    )
+                )
+
+                xmp = build_drawthings_xmp(
+                    json_string,
+                    build_description(metadata),
+                )
+
+                itxt = build_itxt_chunk(
+                    "XML:com.adobe.xmp",
+                    xmp,
+                )
+
+                out.extend(
+                    png_chunk(
+                        b"iTXt",
+                        itxt,
+                    )
+                )
+
+            inserted = True
+
+    return bytes(out)