rs-flint 0.1.7

Run koharu-torch inference graphs from RustScript host functions
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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

"""Patch koharu-diffusion's C enum wrappers for the current target ABI."""

from __future__ import annotations

import json
import subprocess
from pathlib import Path


def koharu_enums_path() -> Path:
    metadata = subprocess.run(
        ["cargo", "metadata", "--locked", "--format-version", "1"],
        check=True,
        stdout=subprocess.PIPE,
        text=True,
        encoding="utf-8",
    )
    packages = json.loads(metadata.stdout)["packages"]
    matches = [
        package
        for package in packages
        if package["name"] == "koharu-diffusion"
        and package["source"].startswith("git+")
    ]
    if len(matches) != 1:
        raise RuntimeError(
            f"expected one Git koharu-diffusion package, found {len(matches)}"
        )
    return Path(matches[0]["manifest_path"]).parent / "src" / "enums.rs"


path = koharu_enums_path()
source = path.read_text(encoding="utf-8")

marker = "// Patched by Flint: preserve each bindgen enum's target-specific ABI."
if marker in source:
    print(f"already patched {path}")
    raise SystemExit(0)

replacements = {
    "use crate::{Error, Result, ffi::NativeCall, sys};\n": (
        "use crate::{Error, Result, ffi::NativeCall, sys};\n\n" + marker + "\n"
    ),
    "pub enum $name:ident, $kind:literal, $name_fn:path, $parse_fn:path, $invalid:path;": (
        "pub enum $name:ident, $kind:literal, $raw_ty:ty, $name_fn:path, "
        "$parse_fn:path, $invalid:path;"
    ),
    "pub enum $name:ident, $kind:literal {": (
        "pub enum $name:ident, $kind:literal, $raw_ty:ty {"
    ),
    "$($variant = $raw),+": "$($variant = $raw as i32),+",
    "pub const fn as_raw(self) -> i32": "pub const fn as_raw(self) -> $raw_ty",
    "self as i32": "self as i32 as $raw_ty",
    "impl TryFrom<i32> for $name": "impl TryFrom<$raw_ty> for $name",
    "fn try_from(value: i32) -> Result<Self>": "fn try_from(value: $raw_ty) -> Result<Self>",
    "value => Err(Error::InvalidEnum { kind: $kind, value }),": (
        "value => Err(Error::InvalidEnum { kind: $kind, value: value as i32 }),"
    ),
    'pub enum WeightType, "weight type",': (
        'pub enum WeightType, "weight type", sys::sd_type_t,'
    ),
    'pub enum RngType, "RNG type",': (
        'pub enum RngType, "RNG type", sys::rng_type_t,'
    ),
    'pub enum SampleMethod, "sample method",': (
        'pub enum SampleMethod, "sample method", sys::sample_method_t,'
    ),
    'pub enum Scheduler, "scheduler",': (
        'pub enum Scheduler, "scheduler", sys::scheduler_t,'
    ),
    'pub enum Prediction, "prediction",': (
        'pub enum Prediction, "prediction", sys::prediction_t,'
    ),
    'pub enum PreviewMode, "preview mode",': (
        'pub enum PreviewMode, "preview mode", sys::preview_t,'
    ),
    'pub enum LoraApplyMode, "LoRA apply mode",': (
        'pub enum LoraApplyMode, "LoRA apply mode", sys::lora_apply_mode_t,'
    ),
    'pub enum HiresUpscaler, "high-resolution upscaler",': (
        'pub enum HiresUpscaler, "high-resolution upscaler", '
        'sys::sd_hires_upscaler_t,'
    ),
    'pub enum VaeFormat, "VAE format" {': (
        'pub enum VaeFormat, "VAE format", sys::sd_vae_format_t {'
    ),
    'pub enum CacheMode, "cache mode" {': (
        'pub enum CacheMode, "cache mode", sys::sd_cache_mode_t {'
    ),
    'pub enum LogLevel, "log level" {': (
        'pub enum LogLevel, "log level", sys::sd_log_level_t {'
    ),
    'pub enum CancelMode, "cancel mode" {': (
        'pub enum CancelMode, "cancel mode", sys::sd_cancel_mode_t {'
    ),
}

expected_counts = {
    "use crate::{Error, Result, ffi::NativeCall, sys};\n": 1,
    "pub enum $name:ident, $kind:literal, $name_fn:path, $parse_fn:path, $invalid:path;": 1,
    "pub enum $name:ident, $kind:literal {": 1,
    "$($variant = $raw),+": 2,
    "pub const fn as_raw(self) -> i32": 2,
    "self as i32": 2,
    "impl TryFrom<i32> for $name": 2,
    "fn try_from(value: i32) -> Result<Self>": 2,
    "value => Err(Error::InvalidEnum { kind: $kind, value }),": 2,
    'pub enum WeightType, "weight type",': 1,
    'pub enum RngType, "RNG type",': 1,
    'pub enum SampleMethod, "sample method",': 1,
    'pub enum Scheduler, "scheduler",': 1,
    'pub enum Prediction, "prediction",': 1,
    'pub enum PreviewMode, "preview mode",': 1,
    'pub enum LoraApplyMode, "LoRA apply mode",': 1,
    'pub enum HiresUpscaler, "high-resolution upscaler",': 1,
    'pub enum VaeFormat, "VAE format" {': 1,
    'pub enum CacheMode, "cache mode" {': 1,
    'pub enum LogLevel, "log level" {': 1,
    'pub enum CancelMode, "cancel mode" {': 1,
}

for old, new in replacements.items():
    count = source.count(old)
    if count != expected_counts[old]:
        raise RuntimeError(f"unexpected occurrence count for {old!r}: {count}")
    source = source.replace(old, new)

path.write_text(source, encoding="utf-8")
print(f"patched {path}")