Add image handling and denoise parameter support in inference and model components

Author: Xayah-Hina

.github/workflows/docker.yml

@@ -29,6 +29,11 @@ jobs:
         with:
           persist-credentials: false
 
+      - name: Free disk space for CUDA layers
+        run: |
+          sudo rm --recursive --force /usr/share/dotnet /usr/local/lib/android /opt/hostedtoolcache
+          df -h /
+
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@37fe631027851001ddb9b187196cc803df7f5f0e # v4.3.0
 
@@ -59,8 +64,7 @@ jobs:
           tags: ${{ steps.metadata.outputs.tags }}
           labels: ${{ steps.metadata.outputs.labels }}
           annotations: ${{ steps.metadata.outputs.annotations }}
-          load: ${{ env.PUBLISH != 'true' }}
-          push: ${{ env.PUBLISH == 'true' }}
+          outputs: ${{ env.PUBLISH == 'true' && 'type=image,push=true' || 'type=cacheonly' }}
           provenance: ${{ env.PUBLISH == 'true' }}
           sbom: ${{ env.PUBLISH == 'true' }}
           cache-from: type=gha,scope=genesia-docker

CMakeLists.txt

@@ -127,9 +127,10 @@ target_link_libraries(
         PUBLIC
         CCCL::libcudacxx
         CUDA::toolkit
-        genesia::cudnn-frontend
+        cuDNN::cuDNN
         CUDA::cublasLt
         PRIVATE
+        genesia::cudnn-frontend
         CUDA::cudart
         CUDA::cuda_driver
         CUDA::nvrtc

Dockerfile

@@ -1,21 +1,25 @@
 # syntax=docker/dockerfile:1
 # hadolint global ignore=DL3007
 
-FROM archlinux:latest AS build
+FROM archlinux:latest AS cuda
 
 SHELL ["/bin/bash", "-euo", "pipefail", "-c"]
 
 RUN --mount=type=cache,target=/var/cache/pacman/pkg,sharing=locked \
-    pacman -Syu --noconfirm --needed \
+    pacman -Syu --noconfirm --needed cuda cudnn gcc-libs
+
+ENV PATH="/opt/cuda/bin:${PATH}"
+
+
+FROM cuda AS build
+
+RUN --mount=type=cache,target=/var/cache/pacman/pkg,sharing=locked \
+    pacman -S --noconfirm --needed \
         base-devel \
         cmake \
-        cuda \
-        cudnn \
         git \
         ninja
 
-ENV PATH="/opt/cuda/bin:${PATH}"
-
 WORKDIR /src
 COPY --link . .
 
@@ -30,11 +34,9 @@ RUN cmake -S . -B cmake-build-release -G Ninja \
     && cmake --build cmake-build-release --target genesia --parallel
 
 
-FROM archlinux:latest AS runtime
+FROM cuda AS runtime
 
-RUN --mount=type=cache,target=/var/cache/pacman/pkg,sharing=locked \
-    pacman -Syu --noconfirm --needed cuda cudnn gcc-libs \
-    && groupadd --gid 10001 genesia \
+RUN groupadd --gid 10001 genesia \
     && useradd --uid 10001 --gid 10001 --home-dir /workspace --shell /usr/bin/nologin genesia \
     && install --directory --owner=10001 --group=10001 /opt/genesia/bin /workspace
 

genesia/core/generation/defaults.ixx

@@ -6,6 +6,7 @@ export namespace genesia::defaults {
     inline constexpr int height       = 1536;
     inline constexpr int steps        = 50;
     inline constexpr float cfg        = 4.5F;
+    inline constexpr float denoise    = 0.50F;
     inline constexpr bool random_seed = true;
     inline constexpr std::array<std::uint64_t, 1> seeds{16494404764960740964ULL};
     inline constexpr int warmup              = 0;

genesia/core/generation/output.cpp

@@ -49,6 +49,10 @@ namespace genesia {
         const auto started = std::chrono::steady_clock::now();
         const auto model   = record.model.u8string();
         nlohmann::json metadata{{"version", 1}, {"model", std::string{model.begin(), model.end()}}, {"seed", record.seed}, {"steps", record.parameters.steps}, {"cfg", record.parameters.cfg}, {"sampler", "euler"}, {"scheduler", "simple"}};
+        if (!record.source.empty()) {
+            const auto source = record.source.u8string();
+            metadata["repaint"] = {{"source", std::string{source.begin(), source.end()}}, {"denoise", record.parameters.denoise}};
+        }
         for (const auto& [name, side, text] : {std::tuple{"positive", &record.prompt.positive, &record.parameters.positive}, std::tuple{"negative", &record.prompt.negative, &record.parameters.negative}}) {
             auto& saved     = metadata["prompt"][name];
             saved["text"]   = *text;

genesia/core/generation/output.ixx

@@ -12,6 +12,7 @@ export namespace genesia {
         std::filesystem::path model;
         prompt::Pair prompt;
         std::shared_ptr<const prompt::Catalog> catalog;
+        std::filesystem::path source;
     };
     struct ImageWriter final {
         std::filesystem::path directory;

genesia/core/neural/inference-runtime.cpp

@@ -12,6 +12,25 @@ module genesia.neural.inference_runtime;
 import std;
 
 namespace genesia::neural {
+    struct InferenceRuntime::ConvPlan final {
+        std::array<int, 10> key;
+        cudnn_frontend::graph::Graph graph;
+
+        explicit ConvPlan(const std::array<int, 10>& key);
+    };
+
+    struct InferenceRuntime::AttentionPlan final {
+        std::array<int, 10> key;
+        cudnn_frontend::graph::Graph graph;
+        ::cuda::device_buffer<std::int32_t> query_lengths;
+
+        AttentionPlan(::cuda::stream_ref stream, const std::array<int, 10>& shape);
+    };
+
+    void check(cudnn_frontend::error_t status) {
+        if (status.is_bad()) throw std::runtime_error{status.get_message()};
+    }
+
     nlohmann::json read_plan(const std::filesystem::path& path) {
         std::ifstream file{path, std::ios::binary};
         file.exceptions(std::ios::failbit | std::ios::badbit);
@@ -50,7 +69,7 @@ namespace genesia::neural {
         check(cublasLtMatrixLayoutCreate(std::out_ptr(output), result_dtype, n, m, n));
     }
 
-    ConvPlan::ConvPlan(const std::array<int, 10>& shape) : key{shape} {
+    InferenceRuntime::ConvPlan::ConvPlan(const std::array<int, 10>& shape) : key{shape} {
         const auto [n, h, w, input_width, output_width, kernel, stride, padding, scalar, residual] = shape;
         const auto dtype                                                                           = scalar == 1 ? cudnn_frontend::DataType_t::HALF : cudnn_frontend::DataType_t::BFLOAT16;
         graph.set_io_data_type(dtype).set_intermediate_data_type(cudnn_frontend::DataType_t::FLOAT).set_compute_data_type(cudnn_frontend::DataType_t::FLOAT);
@@ -320,7 +339,4 @@ namespace genesia::neural {
     void check(const cudnnStatus_t status) {
         if (status != CUDNN_STATUS_SUCCESS) throw std::runtime_error{cudnnGetErrorString(status)};
     }
-    void check(cudnn_frontend::error_t status) {
-        if (status.is_bad()) throw std::runtime_error{status.get_message()};
-    }
 } // namespace genesia::neural

genesia/core/neural/inference-runtime.ixx

@@ -1,7 +1,7 @@
 module;
 
 #include <cublasLt.h>
-#include <cudnn_frontend.h>
+#include <cudnn.h>
 #include <genesia/cuda.h>
 
 export module genesia.neural.inference_runtime;
@@ -69,19 +69,15 @@ export namespace genesia::neural {
         MatmulPlan& operator=(const MatmulPlan&) = delete;
     };
 
-    struct ConvPlan final {
-        std::array<int, 10> key;
-        cudnn_frontend::graph::Graph graph;
-
-        explicit ConvPlan(const std::array<int, 10>& key);
-    };
-
     struct InferenceRuntime final {
         ::cuda::stream_ref stream;
         std::size_t cache_hits{};
         std::size_t cache_misses{};
 
     private:
+        struct ConvPlan;
+        struct AttentionPlan;
+
         std::unique_ptr<std::remove_pointer_t<cublasLtHandle_t>, decltype(&cublasLtDestroy)> blas{nullptr, cublasLtDestroy};
         std::unique_ptr<std::remove_pointer_t<cudnnHandle_t>, decltype(&cudnnDestroy)> dnn{nullptr, cudnnDestroy};
         ::cuda::device_buffer<std::byte> workspace;
@@ -94,12 +90,6 @@ export namespace genesia::neural {
         std::size_t intermediate_bytes{};
         std::list<MatmulPlan> matmuls;
         std::list<ConvPlan> convolutions;
-        struct AttentionPlan final {
-            std::array<int, 10> key;
-            cudnn_frontend::graph::Graph graph;
-            ::cuda::device_buffer<std::int32_t> query_lengths;
-            AttentionPlan(::cuda::stream_ref stream, const std::array<int, 10>& shape);
-        };
         std::list<AttentionPlan> attentions;
 
     public:
@@ -124,5 +114,4 @@ export namespace genesia::neural {
     void check(cudaError_t status);
     void check(cublasStatus_t status);
     void check(cudnnStatus_t status);
-    void check(cudnn_frontend::error_t status);
 } // namespace genesia::neural

genesia/core/sdxl/kernels.cu

@@ -71,16 +71,18 @@ namespace genesia::sdxl::kernels {
             if (threadIdx.x * 4 + j < 1000) training[threadIdx.x * 4 + j] = sqrtf(expm1f(-values[j]));
     }
 
-    __global__ void schedule_kernel(SamplingStep* output, float* times, const float* training, const int steps) {
+    __global__ void schedule_kernel(SamplingStep* output, float* times, const float* training, const int steps, const int total) {
         const int i = blockIdx.x * blockDim.x + threadIdx.x;
         if (i > steps) return;
-        const float sigma = i == steps ? 0.0F : training[999 - i * 1000 / steps];
-        const float next  = i + 1 >= steps ? 0.0F : training[999 - (i + 1) * 1000 / steps];
+        const int position = total - steps + i;
+        const int time = 999 - int(static_cast<long long>(position) * 1000 / total);
+        const float sigma = i == steps ? 0.0F : training[time];
+        const float next = i + 1 >= steps ? 0.0F : training[999 - int(static_cast<long long>(position + 1) * 1000 / total)];
         output[i]         = {sigma, next - sigma, rsqrtf(fmaf(sigma, sigma, 1.0F))};
-        if (i < steps) times[i] = float(999 - i * 1000 / steps);
+        if (i < steps) times[i] = float(time);
     }
 
-    __global__ void initialize_kernel(float* state, __half* input, int* step, const std::uint64_t* seed, const SamplingStep* schedule, const int count) {
+    __global__ void initialize_kernel(float* state, __half* input, int* step, const std::uint64_t* seed, const SamplingStep* schedule, const int count, const float* source, const bool full_noise) {
         const int i = blockIdx.x * blockDim.x + threadIdx.x;
         if (i >= count / 4) return;
         uint4 counter{static_cast<unsigned>(i), 0, 0, 0};
@@ -98,8 +100,15 @@ namespace genesia::sdxl::kernels {
         sincosf(float(counter.w) * 0x1p-32F * 6.283185307179586F, &sine1, &cosine1);
         const float radius0 = sqrtf(-2.0F * logf((float(counter.x) + 1.0F) * 0x1p-32F));
         const float radius1 = sqrtf(-2.0F * logf((float(counter.z) + 1.0F) * 0x1p-32F));
-        const float scale   = sqrtf(fmaf(schedule[0].sigma, schedule[0].sigma, 1.0F));
-        const float4 values{radius0 * sine0 * scale, radius0 * cosine0 * scale, radius1 * sine1 * scale, radius1 * cosine1 * scale};
+        const float scale = full_noise ? sqrtf(fmaf(schedule[0].sigma, schedule[0].sigma, 1.0F)) : schedule[0].sigma;
+        float4 values{radius0 * sine0 * scale, radius0 * cosine0 * scale, radius1 * sine1 * scale, radius1 * cosine1 * scale};
+        if (source) {
+            const float4 original = reinterpret_cast<const float4*>(source)[i];
+            values.x += original.x;
+            values.y += original.y;
+            values.z += original.z;
+            values.w += original.w;
+        }
         reinterpret_cast<float4*>(state)[i] = values;
         const float data[4]{values.x, values.y, values.z, values.w};
 #pragma unroll
@@ -148,6 +157,24 @@ namespace genesia::sdxl::kernels {
         if (stopped || completed == count) ::cuda::atomic_ref<std::uint32_t, ::cuda::thread_scope_system>{control->stage}.store(static_cast<std::uint32_t>(stopped ? Stage::cancelled : Stage::decoding), ::cuda::memory_order_release);
     }
 
+    __global__ void image_encode_kernel(__nv_bfloat16* output, const std::uint8_t* input, const int count) {
+        const int i = blockIdx.x * blockDim.x + threadIdx.x;
+        if (i < count) output[i] = __nv_bfloat16(float(input[i]) * (2.0F / 255.0F) - 1.0F);
+    }
+
+    __global__ void encoder_pad_kernel(__nv_bfloat16* output, const __nv_bfloat16* input, const int height, const int width, const int channels) {
+        const int i = blockIdx.x * blockDim.x + threadIdx.x;
+        if (i >= (height + 1) * (width + 1) * channels) return;
+        const int x = i / channels % (width + 1);
+        const int y = i / channels / (width + 1);
+        output[i] = x < width && y < height ? input[(y * width + x) * channels + i % channels] : __nv_bfloat16(0.0F);
+    }
+
+    __global__ void latent_encode_kernel(float* output, const __nv_bfloat16* moments, const int count) {
+        const int i = blockIdx.x * blockDim.x + threadIdx.x;
+        if (i < count) output[i] = float(moments[i / 4 * 8 + i % 4]) * 0.13025F;
+    }
+
     __global__ void latent_decode_kernel(__nv_bfloat16* output, const float* input, const int count) {
         const int i = blockIdx.x * blockDim.x + threadIdx.x;
         if (i < count) output[i] = __nv_bfloat16(input[i] / 0.13025F);
@@ -184,11 +211,11 @@ namespace genesia::sdxl::kernels {
     void training_sigmas(const ::cuda::stream_ref stream, float* output) {
         ::cuda::launch(stream, ::cuda::make_config(::cuda::make_hierarchy(::cuda::grid_dims(1), ::cuda::block_dims(256))), training_kernel, output);
     }
-    void prepare_schedule(const ::cuda::stream_ref stream, SamplingStep* output, float* times, const float* training, const int steps) {
-        ::cuda::launch(stream, ::cuda::make_config(::cuda::make_hierarchy(::cuda::grid_dims((steps + 256) / 256), ::cuda::block_dims(256))), schedule_kernel, output, times, training, steps);
+    void prepare_schedule(const ::cuda::stream_ref stream, SamplingStep* output, float* times, const float* training, const int steps, const float denoise) {
+        ::cuda::launch(stream, ::cuda::make_config(::cuda::make_hierarchy(::cuda::grid_dims((steps + 256) / 256), ::cuda::block_dims(256))), schedule_kernel, output, times, training, steps, static_cast<int>(steps / double(denoise)));
     }
-    void initialize(const ::cuda::stream_ref stream, float* state, void* input, int* step, const std::uint64_t* seed, const SamplingStep* schedule, const int count) {
-        ::cuda::launch(stream, ::cuda::make_config(::cuda::make_hierarchy(::cuda::grid_dims((count / 4 + 255) / 256), ::cuda::block_dims(256))), initialize_kernel, state, static_cast<__half*>(input), step, seed, schedule, count);
+    void initialize(const ::cuda::stream_ref stream, float* state, void* input, int* step, const std::uint64_t* seed, const SamplingStep* schedule, const int count, const float* source, const bool full_noise) {
+        ::cuda::launch(stream, ::cuda::make_config(::cuda::make_hierarchy(::cuda::grid_dims((count / 4 + 255) / 256), ::cuda::block_dims(256))), initialize_kernel, state, static_cast<__half*>(input), step, seed, schedule, count, source, full_noise);
     }
     void snapshot_begin(const ::cuda::stream_ref stream, int* selected, SnapshotSlot* slots) {
         ::cuda::launch(stream, ::cuda::make_config(::cuda::make_hierarchy(::cuda::grid_dims(1), ::cuda::block_dims(1))), snapshot_begin_kernel, selected, slots);
@@ -206,6 +233,15 @@ namespace genesia::sdxl::kernels {
     void advance(const ::cuda::stream_ref stream, int* step, const int count, const cudaGraphConditionalHandle loop, const cudaGraphConditionalHandle decode, Control* control) {
         ::cuda::launch(stream, ::cuda::make_config(::cuda::make_hierarchy(::cuda::grid_dims(1), ::cuda::block_dims(1))), advance_kernel, step, count, loop, decode, control);
     }
+    void image_encode(const ::cuda::stream_ref stream, void* output, const std::uint8_t* pixels, const int count) {
+        ::cuda::launch(stream, ::cuda::make_config(::cuda::make_hierarchy(::cuda::grid_dims((count + 255) / 256), ::cuda::block_dims(256))), image_encode_kernel, static_cast<__nv_bfloat16*>(output), pixels, count);
+    }
+    void encoder_pad(const ::cuda::stream_ref stream, void* output, const void* input, const int height, const int width, const int channels) {
+        ::cuda::launch(stream, ::cuda::make_config(::cuda::make_hierarchy(::cuda::grid_dims(((height + 1) * (width + 1) * channels + 255) / 256), ::cuda::block_dims(256))), encoder_pad_kernel, static_cast<__nv_bfloat16*>(output), static_cast<const __nv_bfloat16*>(input), height, width, channels);
+    }
+    void latent_encode(const ::cuda::stream_ref stream, float* output, const void* moments, const int count) {
+        ::cuda::launch(stream, ::cuda::make_config(::cuda::make_hierarchy(::cuda::grid_dims((count + 255) / 256), ::cuda::block_dims(256))), latent_encode_kernel, output, static_cast<const __nv_bfloat16*>(moments), count);
+    }
     void latent_decode(const ::cuda::stream_ref stream, void* output, const float* latent, const int count) {
         ::cuda::launch(stream, ::cuda::make_config(::cuda::make_hierarchy(::cuda::grid_dims((count + 255) / 256), ::cuda::block_dims(256))), latent_decode_kernel, static_cast<__nv_bfloat16*>(output), latent, count);
     }

genesia/core/sdxl/kernels.h

@@ -20,12 +20,15 @@ namespace genesia::sdxl::kernels {
     void conditions(const ::cuda::stream_ref stream, void* output, const void* pooled, const void* geometry);
     void time_condition(const ::cuda::stream_ref stream, void* output, const void* times, const void* labels, const int steps);
     void training_sigmas(const ::cuda::stream_ref stream, float* output);
-    void prepare_schedule(const ::cuda::stream_ref stream, SamplingStep* output, float* times, const float* training, const int steps);
-    void initialize(const ::cuda::stream_ref stream, float* state, void* input, int* step, const std::uint64_t* seed, const SamplingStep* schedule, const int count);
+    void prepare_schedule(::cuda::stream_ref stream, SamplingStep* output, float* times, const float* training, int steps, float denoise = 1);
+    void initialize(::cuda::stream_ref stream, float* state, void* input, int* step, const std::uint64_t* seed, const SamplingStep* schedule, int count, const float* source = nullptr, bool full_noise = true);
     void snapshot_begin(::cuda::stream_ref stream, int* selected, SnapshotSlot* slots);
     void snapshot_publish(::cuda::stream_ref stream, const int* selected, SnapshotSlot* slots, const int* step);
     void euler(const ::cuda::stream_ref stream, float* state, void* input, const void* epsilon, const SamplingStep* schedule, const int* step, const float cfg, const int count, float* snapshots = nullptr, const int* selected = nullptr);
     void advance(::cuda::stream_ref stream, int* step, int count, cudaGraphConditionalHandle loop, cudaGraphConditionalHandle decode, Control* control);
+    void image_encode(::cuda::stream_ref stream, void* output, const std::uint8_t* pixels, int count);
+    void encoder_pad(::cuda::stream_ref stream, void* output, const void* input, int height, int width, int channels);
+    void latent_encode(::cuda::stream_ref stream, float* output, const void* moments, int count);
     void latent_decode(const ::cuda::stream_ref stream, void* output, const float* latent, const int count);
     void pixels(const ::cuda::stream_ref stream, std::uint8_t* output, const void* input, const int count);
 } // namespace genesia::sdxl::kernels