Duplicate from lnyan/stablediffusion-infinity

Co-authored-by: Lnyan <lnyan@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,31 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

PyPatchMatch/.gitignore

@@ -0,0 +1,4 @@
+/build/
+/*.so
+__pycache__
+*.py[cod]

PyPatchMatch/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 Jiayuan Mao
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

PyPatchMatch/Makefile

@@ -0,0 +1,54 @@
+#
+# Makefile
+# Jiayuan Mao, 2019-01-09 13:59
+#
+
+SRC_DIR = csrc
+INC_DIR = csrc
+OBJ_DIR = build/obj
+TARGET = libpatchmatch.so
+
+LIB_TARGET = $(TARGET)
+INCLUDE_DIR = -I $(SRC_DIR) -I $(INC_DIR)
+
+CXX = $(ENVIRONMENT_OPTIONS) g++
+CXXFLAGS = -std=c++14
+CXXFLAGS += -Ofast -ffast-math -w
+# CXXFLAGS += -g
+CXXFLAGS += $(shell pkg-config --cflags opencv) -fPIC
+CXXFLAGS += $(INCLUDE_DIR)
+LDFLAGS = $(shell pkg-config --cflags --libs opencv) -shared -fPIC
+
+
+CXXSOURCES = $(shell find $(SRC_DIR)/ -name "*.cpp")
+OBJS = $(addprefix $(OBJ_DIR)/,$(CXXSOURCES:.cpp=.o))
+DEPFILES = $(OBJS:.o=.d)
+
+.PHONY: all clean rebuild test
+
+all: $(LIB_TARGET)
+
+$(OBJ_DIR)/%.o: %.cpp
+	@echo "[CC] $< ..."
+	@$(CXX) -c $< $(CXXFLAGS) -o $@
+
+$(OBJ_DIR)/%.d: %.cpp
+	@mkdir -pv $(dir $@)
+	@echo "[dep] $< ..."
+	@$(CXX) $(INCLUDE_DIR) $(CXXFLAGS) -MM -MT "$(OBJ_DIR)/$(<:.cpp=.o) $(OBJ_DIR)/$(<:.cpp=.d)" "$<" > "$@"
+
+sinclude $(DEPFILES)
+
+$(LIB_TARGET): $(OBJS)
+	@echo "[link] $(LIB_TARGET) ..."
+	@$(CXX) $(OBJS) -o $@ $(CXXFLAGS) $(LDFLAGS)
+
+clean:
+	rm -rf $(OBJ_DIR) $(LIB_TARGET)
+
+rebuild:
+	+@make clean
+	+@make
+
+# vim:ft=make
+#

PyPatchMatch/README.md

@@ -0,0 +1,64 @@
+PatchMatch based Inpainting
+=====================================
+This library implements the PatchMatch based inpainting algorithm. It provides both C++ and Python interfaces.
+This implementation is heavily based on the implementation by Younesse ANDAM:
+(younesse-cv/PatchMatch)[https://github.com/younesse-cv/PatchMatch], with some bugs fix.
+
+Usage
+-------------------------------------
+
+You need to first install OpenCV to compile the C++ libraries. Then, run `make` to compile the
+shared library `libpatchmatch.so`.
+
+For Python users (example available at `examples/py_example.py`)
+
+```python
+import patch_match
+
+image = ...  # either a numpy ndarray or a PIL Image object.
+mask = ...   # either a numpy ndarray or a PIL Image object.
+result = patch_match.inpaint(image, mask, patch_size=5)
+```
+
+For C++ users (examples available at `examples/cpp_example.cpp`)
+
+```cpp
+#include "inpaint.h"
+
+int main() {
+    cv::Mat image = ...
+    cv::Mat mask = ...
+
+    cv::Mat result = Inpainting(image, mask, 5).run();
+
+    return 0;
+}
+```
+
+
+README and COPYRIGHT by Younesse ANDAM
+-------------------------------------
+@Author: Younesse ANDAM
+
+@Contact: younesse.andam@gmail.com
+
+Description: This project is a personal implementation of an algorithm called PATCHMATCH that restores missing areas in an image.
+The algorithm is presented in the following paper
+ PatchMatch  A Randomized Correspondence Algorithm
+               for Structural Image Editing
+   by C.Barnes,E.Shechtman,A.Finkelstein and Dan B.Goldman
+   ACM Transactions on Graphics (Proc. SIGGRAPH), vol.28, aug-2009
+
+ For more information please refer to
+ http://www.cs.princeton.edu/gfx/pubs/Barnes_2009_PAR/index.php
+
+Copyright (c) 2010-2011
+
+
+Requirements
+-------------------------------------
+
+To run the project you need to install Opencv library and link it to your project.
+Opencv can be download it here
+http://opencv.org/downloads.html
+

PyPatchMatch/csrc/inpaint.cpp

@@ -0,0 +1,234 @@
+#include <algorithm>
+#include <iostream>
+#include <opencv2/imgcodecs.hpp>
+#include <opencv2/imgproc.hpp>
+#include <opencv2/highgui.hpp>
+
+#include "inpaint.h"
+
+namespace {
+    static std::vector<double> kDistance2Similarity;
+
+    void init_kDistance2Similarity() {
+        double base[11] = {1.0, 0.99, 0.96, 0.83, 0.38, 0.11, 0.02, 0.005, 0.0006, 0.0001, 0};
+        int length = (PatchDistanceMetric::kDistanceScale + 1);
+        kDistance2Similarity.resize(length);
+        for (int i = 0; i < length; ++i) {
+            double t = (double) i / length;
+            int j = (int) (100 * t);
+            int k = j + 1;
+            double vj = (j < 11) ? base[j] : 0;
+            double vk = (k < 11) ? base[k] : 0;
+            kDistance2Similarity[i] = vj + (100 * t - j) * (vk - vj);
+        }
+    }
+
+
+    inline void _weighted_copy(const MaskedImage &source, int ys, int xs, cv::Mat &target, int yt, int xt, double weight) {
+        if (source.is_masked(ys, xs)) return;
+        if (source.is_globally_masked(ys, xs)) return;
+
+        auto source_ptr = source.get_image(ys, xs);
+        auto target_ptr = target.ptr<double>(yt, xt);
+
+#pragma unroll
+        for (int c = 0; c < 3; ++c)
+            target_ptr[c] += static_cast<double>(source_ptr[c]) * weight;
+        target_ptr[3] += weight;
+    }
+}
+
+/**
+ * This algorithme uses a version proposed by Xavier Philippeau.
+ */
+
+Inpainting::Inpainting(cv::Mat image, cv::Mat mask, const PatchDistanceMetric *metric)
+    : m_initial(image, mask), m_distance_metric(metric), m_pyramid(), m_source2target(), m_target2source() {
+    _initialize_pyramid();
+}
+
+Inpainting::Inpainting(cv::Mat image, cv::Mat mask, cv::Mat global_mask, const PatchDistanceMetric *metric)
+    : m_initial(image, mask, global_mask), m_distance_metric(metric), m_pyramid(), m_source2target(), m_target2source() {
+    _initialize_pyramid();
+}
+
+void Inpainting::_initialize_pyramid() {
+    auto source = m_initial;
+    m_pyramid.push_back(source);
+    while (source.size().height > m_distance_metric->patch_size() && source.size().width > m_distance_metric->patch_size()) {
+        source = source.downsample();
+        m_pyramid.push_back(source);
+    }
+
+    if (kDistance2Similarity.size() == 0) {
+        init_kDistance2Similarity();
+    }
+}
+
+cv::Mat Inpainting::run(bool verbose, bool verbose_visualize, unsigned int random_seed) {
+    srand(random_seed);
+    const int nr_levels = m_pyramid.size();
+
+    MaskedImage source, target;
+    for (int level = nr_levels - 1; level >= 0; --level) {
+        if (verbose) std::cerr << "Inpainting level: " << level << std::endl;
+
+        source = m_pyramid[level];
+
+        if (level == nr_levels - 1) {
+            target = source.clone();
+            target.clear_mask();
+            m_source2target = NearestNeighborField(source, target, m_distance_metric);
+            m_target2source = NearestNeighborField(target, source, m_distance_metric);
+        } else {
+            m_source2target = NearestNeighborField(source, target, m_distance_metric, m_source2target);
+            m_target2source = NearestNeighborField(target, source, m_distance_metric, m_target2source);
+        }
+
+        if (verbose) std::cerr << "Initialization done." << std::endl;
+
+        if (verbose_visualize) {
+            auto visualize_size = m_initial.size();
+            cv::Mat source_visualize(visualize_size, m_initial.image().type());
+            cv::resize(source.image(), source_visualize, visualize_size);
+            cv::imshow("Source", source_visualize);
+            cv::Mat target_visualize(visualize_size, m_initial.image().type());
+            cv::resize(target.image(), target_visualize, visualize_size);
+            cv::imshow("Target", target_visualize);
+            cv::waitKey(0);
+        }
+
+        target = _expectation_maximization(source, target, level, verbose);
+    }
+
+    return target.image();
+}
+
+// EM-Like algorithm (see "PatchMatch" - page 6).
+// Returns a double sized target image (unless level = 0).
+MaskedImage Inpainting::_expectation_maximization(MaskedImage source, MaskedImage target, int level, bool verbose) {
+    const int nr_iters_em = 1 + 2 * level;
+    const int nr_iters_nnf = static_cast<int>(std::min(7, 1 + level));
+    const int patch_size = m_distance_metric->patch_size();
+
+    MaskedImage new_source, new_target;
+
+    for (int iter_em = 0; iter_em < nr_iters_em; ++iter_em) {
+        if (iter_em != 0) {
+            m_source2target.set_target(new_target);
+            m_target2source.set_source(new_target);
+            target = new_target;
+        }
+
+        if (verbose) std::cerr << "EM Iteration: " << iter_em << std::endl;
+
+        auto size = source.size();
+        for (int i = 0; i < size.height; ++i) {
+            for (int j = 0; j < size.width; ++j) {
+                if (!source.contains_mask(i, j, patch_size)) {
+                    m_source2target.set_identity(i, j);
+                    m_target2source.set_identity(i, j);
+                }
+            }
+        }
+        if (verbose) std::cerr << "  NNF minimization started." << std::endl;
+        m_source2target.minimize(nr_iters_nnf);
+        m_target2source.minimize(nr_iters_nnf);
+        if (verbose) std::cerr << "  NNF minimization finished." << std::endl;
+
+        // Instead of upsizing the final target, we build the last target from the next level source image.
+        // Thus, the final target is less blurry (see "Space-Time Video Completion" - page 5).
+        bool upscaled = false;
+        if (level >= 1 && iter_em == nr_iters_em - 1) {
+            new_source = m_pyramid[level - 1];
+            new_target = target.upsample(new_source.size().width, new_source.size().height, m_pyramid[level - 1].global_mask());
+            upscaled = true;
+        } else {
+            new_source = m_pyramid[level];
+            new_target = target.clone();
+        }
+
+        auto vote = cv::Mat(new_target.size(), CV_64FC4);
+        vote.setTo(cv::Scalar::all(0));
+
+        // Votes for best patch from NNF Source->Target (completeness) and Target->Source (coherence).
+        _expectation_step(m_source2target, 1, vote, new_source, upscaled);
+        if (verbose) std::cerr << "  Expectation source to target finished." << std::endl;
+        _expectation_step(m_target2source, 0, vote, new_source, upscaled);
+        if (verbose) std::cerr << "  Expectation target to source finished." << std::endl;
+
+        // Compile votes and update pixel values.
+        _maximization_step(new_target, vote);
+        if (verbose) std::cerr << "  Minimization step finished." << std::endl;
+    }
+
+    return new_target;
+}
+
+// Expectation step: vote for best estimations of each pixel.
+void Inpainting::_expectation_step(
+    const NearestNeighborField &nnf, bool source2target,
+    cv::Mat &vote, const MaskedImage &source, bool upscaled
+) {
+    auto source_size = nnf.source_size();
+    auto target_size = nnf.target_size();
+    const int patch_size = m_distance_metric->patch_size();
+
+    for (int i = 0; i < source_size.height; ++i) {
+        for (int j = 0; j < source_size.width; ++j) {
+            if (nnf.source().is_globally_masked(i, j)) continue;
+            int yp = nnf.at(i, j, 0), xp = nnf.at(i, j, 1), dp = nnf.at(i, j, 2);
+            double w = kDistance2Similarity[dp];
+
+            for (int di = -patch_size; di <= patch_size; ++di) {
+                for (int dj = -patch_size; dj <= patch_size; ++dj) {
+                    int ys = i + di, xs = j + dj, yt = yp + di, xt = xp + dj;
+                    if (!(ys >= 0 && ys < source_size.height && xs >= 0 && xs < source_size.width)) continue;
+                    if (nnf.source().is_globally_masked(ys, xs)) continue;
+                    if (!(yt >= 0 && yt < target_size.height && xt >= 0 && xt < target_size.width)) continue;
+                    if (nnf.target().is_globally_masked(yt, xt)) continue;
+
+                    if (!source2target) {
+                        std::swap(ys, yt);
+                        std::swap(xs, xt);
+                    }
+
+                    if (upscaled) {
+                        for (int uy = 0; uy < 2; ++uy) {
+                            for (int ux = 0; ux < 2; ++ux) {
+                                _weighted_copy(source, 2 * ys + uy, 2 * xs + ux, vote, 2 * yt + uy, 2 * xt + ux, w);
+                            }
+                        }
+                    } else {
+                        _weighted_copy(source, ys, xs, vote, yt, xt, w);
+                    }
+                }
+            }
+        }
+    }
+}
+
+// Maximization Step: maximum likelihood of target pixel.
+void Inpainting::_maximization_step(MaskedImage &target, const cv::Mat &vote) {
+    auto target_size = target.size();
+    for (int i = 0; i < target_size.height; ++i) {
+        for (int j = 0; j < target_size.width; ++j) {
+            const double *source_ptr = vote.ptr<double>(i, j);
+            unsigned char *target_ptr = target.get_mutable_image(i, j);
+
+            if (target.is_globally_masked(i, j)) {
+                continue;
+            }
+
+            if (source_ptr[3] > 0) {
+                unsigned char r = cv::saturate_cast<unsigned char>(source_ptr[0] / source_ptr[3]);
+                unsigned char g = cv::saturate_cast<unsigned char>(source_ptr[1] / source_ptr[3]);
+                unsigned char b = cv::saturate_cast<unsigned char>(source_ptr[2] / source_ptr[3]);
+                target_ptr[0] = r, target_ptr[1] = g, target_ptr[2] = b;
+            } else {
+                target.set_mask(i, j, 0);
+            }
+        }
+    }
+}
+

PyPatchMatch/csrc/inpaint.h

@@ -0,0 +1,27 @@
+#pragma once
+
+#include <vector>
+
+#include "masked_image.h"
+#include "nnf.h"
+
+class Inpainting {
+public:
+    Inpainting(cv::Mat image, cv::Mat mask, const PatchDistanceMetric *metric);
+    Inpainting(cv::Mat image, cv::Mat mask, cv::Mat global_mask, const PatchDistanceMetric *metric);
+    cv::Mat run(bool verbose = false, bool verbose_visualize = false, unsigned int random_seed = 1212);
+
+private:
+    void _initialize_pyramid(void);
+    MaskedImage _expectation_maximization(MaskedImage source, MaskedImage target, int level, bool verbose);
+    void _expectation_step(const NearestNeighborField &nnf, bool source2target, cv::Mat &vote, const MaskedImage &source, bool upscaled);
+    void _maximization_step(MaskedImage &target, const cv::Mat &vote);
+
+    MaskedImage m_initial;
+    std::vector<MaskedImage> m_pyramid;
+
+    NearestNeighborField m_source2target;
+    NearestNeighborField m_target2source;
+    const PatchDistanceMetric *m_distance_metric;
+};
+

PyPatchMatch/csrc/masked_image.cpp

@@ -0,0 +1,138 @@
+#include "masked_image.h"
+#include <algorithm>
+#include <iostream>
+
+const cv::Size MaskedImage::kDownsampleKernelSize = cv::Size(6, 6);
+const int MaskedImage::kDownsampleKernel[6] = {1, 5, 10, 10, 5, 1};
+
+bool MaskedImage::contains_mask(int y, int x, int patch_size) const {
+    auto mask_size = size();
+    for (int dy = -patch_size; dy <= patch_size; ++dy) {
+        for (int dx = -patch_size; dx <= patch_size; ++dx) {
+            int yy = y + dy, xx = x + dx;
+            if (yy >= 0 && yy < mask_size.height && xx >= 0 && xx < mask_size.width) {
+                if (is_masked(yy, xx) && !is_globally_masked(yy, xx)) return true;
+            }
+        }
+    }
+    return false;
+}
+
+MaskedImage MaskedImage::downsample() const {
+    const auto &kernel_size = MaskedImage::kDownsampleKernelSize;
+    const auto &kernel = MaskedImage::kDownsampleKernel;
+
+    const auto size = this->size();
+    const auto new_size = cv::Size(size.width / 2, size.height / 2);
+
+    auto ret = MaskedImage(new_size.width, new_size.height);
+    if (!m_global_mask.empty()) ret.init_global_mask_mat();
+    for (int y = 0; y < size.height - 1; y += 2) {
+        for (int x = 0; x < size.width - 1; x += 2) {
+            int r = 0, g = 0, b = 0, ksum = 0;
+            bool is_gmasked = true;
+
+            for (int dy = -kernel_size.height / 2 + 1; dy <= kernel_size.height / 2; ++dy) {
+                for (int dx = -kernel_size.width / 2 + 1; dx <= kernel_size.width / 2; ++dx) {
+                    int yy = y + dy, xx = x + dx;
+                    if (yy >= 0 && yy < size.height && xx >= 0 && xx < size.width) {
+                        if (!is_globally_masked(yy, xx)) {
+                            is_gmasked = false;
+                        }
+                        if (!is_masked(yy, xx)) {
+                            auto source_ptr = get_image(yy, xx);
+                            int k = kernel[kernel_size.height / 2 - 1 + dy] * kernel[kernel_size.width / 2 - 1 + dx];
+                            r += source_ptr[0] * k, g += source_ptr[1] * k, b += source_ptr[2] * k;
+                            ksum += k;
+                        }
+                    }
+                }
+            }
+
+            if (ksum > 0) r /= ksum, g /= ksum, b /= ksum;
+
+            if (!m_global_mask.empty()) {
+                ret.set_global_mask(y / 2, x / 2, is_gmasked);
+            }
+            if (ksum > 0) {
+                auto target_ptr = ret.get_mutable_image(y / 2, x / 2);
+                target_ptr[0] = r, target_ptr[1] = g, target_ptr[2] = b;
+                ret.set_mask(y / 2, x / 2, 0);
+            } else {
+                ret.set_mask(y / 2, x / 2, 1);
+            }
+        }
+    }
+
+    return ret;
+}
+
+MaskedImage MaskedImage::upsample(int new_w, int new_h) const {
+    const auto size = this->size();
+    auto ret = MaskedImage(new_w, new_h);
+    if (!m_global_mask.empty()) ret.init_global_mask_mat();
+    for (int y = 0; y < new_h; ++y) {
+        for (int x = 0; x < new_w; ++x) {
+            int yy = y * size.height / new_h;
+            int xx = x * size.width / new_w;
+
+            if (is_globally_masked(yy, xx)) {
+                ret.set_global_mask(y, x, 1);
+                ret.set_mask(y, x, 1);
+            } else {
+                if (!m_global_mask.empty()) ret.set_global_mask(y, x, 0);
+
+                if (is_masked(yy, xx)) {
+                    ret.set_mask(y, x, 1);
+                } else {
+                    auto source_ptr = get_image(yy, xx);
+                    auto target_ptr = ret.get_mutable_image(y, x);
+                    for (int c = 0; c < 3; ++c)
+                        target_ptr[c] = source_ptr[c];
+                    ret.set_mask(y, x, 0);
+                }
+            }
+        }
+    }
+
+    return ret;
+}
+
+MaskedImage MaskedImage::upsample(int new_w, int new_h, const cv::Mat &new_global_mask) const {
+    auto ret = upsample(new_w, new_h);
+    ret.set_global_mask_mat(new_global_mask);
+    return ret;
+}
+
+void MaskedImage::compute_image_gradients() {
+    if (m_image_grad_computed) {
+        return;
+    }
+
+    const auto size = m_image.size();
+    m_image_grady = cv::Mat(size, CV_8UC3);
+    m_image_gradx = cv::Mat(size, CV_8UC3);
+    m_image_grady = cv::Scalar::all(0);
+    m_image_gradx = cv::Scalar::all(0);
+
+    for (int i = 1; i < size.height - 1; ++i) {
+        const auto *ptr = m_image.ptr<unsigned char>(i, 0);
+        const auto *ptry1 = m_image.ptr<unsigned char>(i + 1, 0);
+        const auto *ptry2 = m_image.ptr<unsigned char>(i - 1, 0);
+        const auto *ptrx1 = m_image.ptr<unsigned char>(i, 0) + 3;
+        const auto *ptrx2 = m_image.ptr<unsigned char>(i, 0) - 3;
+        auto *mptry = m_image_grady.ptr<unsigned char>(i, 0);
+        auto *mptrx = m_image_gradx.ptr<unsigned char>(i, 0);
+        for (int j = 3; j < size.width * 3 - 3; ++j) {
+            mptry[j] = (ptry1[j] / 2 - ptry2[j] / 2) + 128;
+            mptrx[j] = (ptrx1[j] / 2 - ptrx2[j] / 2) + 128;
+        }
+    }
+
+    m_image_grad_computed = true;
+}
+
+void MaskedImage::compute_image_gradients() const {
+    const_cast<MaskedImage *>(this)->compute_image_gradients();
+}
+

PyPatchMatch/csrc/masked_image.h

@@ -0,0 +1,112 @@
+#pragma once
+
+#include <opencv2/core.hpp>
+
+class MaskedImage {
+public:
+    MaskedImage() : m_image(), m_mask(), m_global_mask(), m_image_grady(), m_image_gradx(), m_image_grad_computed(false) {
+        // pass
+    }
+    MaskedImage(cv::Mat image, cv::Mat mask) : m_image(image), m_mask(mask), m_image_grad_computed(false) {
+        // pass
+    }
+    MaskedImage(cv::Mat image, cv::Mat mask, cv::Mat global_mask) : m_image(image), m_mask(mask), m_global_mask(global_mask), m_image_grad_computed(false) {
+        // pass
+    }
+    MaskedImage(cv::Mat image, cv::Mat mask, cv::Mat global_mask, cv::Mat grady, cv::Mat gradx, bool grad_computed) :
+        m_image(image), m_mask(mask), m_global_mask(global_mask),
+        m_image_grady(grady), m_image_gradx(gradx), m_image_grad_computed(grad_computed) {
+        // pass
+    }
+    MaskedImage(int width, int height) : m_global_mask(), m_image_grady(), m_image_gradx() {
+        m_image = cv::Mat(cv::Size(width, height), CV_8UC3);
+        m_image = cv::Scalar::all(0);
+
+        m_mask = cv::Mat(cv::Size(width, height), CV_8U);
+        m_mask = cv::Scalar::all(0);
+    }
+    inline MaskedImage clone() {
+        return MaskedImage(
+            m_image.clone(), m_mask.clone(), m_global_mask.clone(),
+            m_image_grady.clone(), m_image_gradx.clone(), m_image_grad_computed
+        );
+    }
+
+    inline cv::Size size() const {
+        return m_image.size();
+    }
+    inline const cv::Mat &image() const {
+        return m_image;
+    }
+    inline const cv::Mat &mask() const {
+        return m_mask;
+    }
+    inline const cv::Mat &global_mask() const {
+        return m_global_mask;
+    }
+    inline const cv::Mat &grady() const {
+        assert(m_image_grad_computed);
+        return m_image_grady;
+    }
+    inline const cv::Mat &gradx() const {
+        assert(m_image_grad_computed);
+        return m_image_gradx;
+    }
+
+    inline void init_global_mask_mat() {
+        m_global_mask = cv::Mat(m_mask.size(), CV_8U);
+        m_global_mask.setTo(cv::Scalar(0));
+    }
+    inline void set_global_mask_mat(const cv::Mat &other) {
+        m_global_mask = other;
+    }
+
+    inline bool is_masked(int y, int x) const {
+        return static_cast<bool>(m_mask.at<unsigned char>(y, x));
+    }
+    inline bool is_globally_masked(int y, int x) const {
+        return !m_global_mask.empty() && static_cast<bool>(m_global_mask.at<unsigned char>(y, x));
+    }
+    inline void set_mask(int y, int x, bool value) {
+        m_mask.at<unsigned char>(y, x) = static_cast<unsigned char>(value);
+    }
+    inline void set_global_mask(int y, int x, bool value) {
+        m_global_mask.at<unsigned char>(y, x) = static_cast<unsigned char>(value);
+    }
+    inline void clear_mask() {
+        m_mask.setTo(cv::Scalar(0));
+    }
+
+    inline const unsigned char *get_image(int y, int x) const {
+        return m_image.ptr<unsigned char>(y, x);
+    }
+    inline unsigned char *get_mutable_image(int y, int x) {
+        return m_image.ptr<unsigned char>(y, x);
+    }
+
+    inline unsigned char get_image(int y, int x, int c) const {
+        return m_image.ptr<unsigned char>(y, x)[c];
+    }
+    inline int get_image_int(int y, int x, int c) const {
+        return static_cast<int>(m_image.ptr<unsigned char>(y, x)[c]);
+    }
+
+    bool contains_mask(int y, int x, int patch_size) const;
+    MaskedImage downsample() const;
+    MaskedImage upsample(int new_w, int new_h) const;
+    MaskedImage upsample(int new_w, int new_h, const cv::Mat &new_global_mask) const;
+    void compute_image_gradients();
+    void compute_image_gradients() const;
+
+    static const cv::Size kDownsampleKernelSize;
+    static const int kDownsampleKernel[6];
+
+private:
+	cv::Mat m_image;
+	cv::Mat m_mask;
+    cv::Mat m_global_mask;
+    cv::Mat m_image_grady;
+    cv::Mat m_image_gradx;
+    bool m_image_grad_computed = false;
+};
+

PyPatchMatch/csrc/nnf.cpp

@@ -0,0 +1,268 @@
+#include <algorithm>
+#include <iostream>
+#include <cmath>
+
+#include "masked_image.h"
+#include "nnf.h"
+
+/**
+* Nearest-Neighbor Field (see PatchMatch algorithm).
+* This algorithme uses a version proposed by Xavier Philippeau.
+*
+*/
+
+template <typename T>
+T clamp(T value, T min_value, T max_value) {
+    return std::min(std::max(value, min_value), max_value);
+}
+
+void NearestNeighborField::_randomize_field(int max_retry, bool reset) {
+    auto this_size = source_size();
+    for (int i = 0; i < this_size.height; ++i) {
+        for (int j = 0; j < this_size.width; ++j) {
+            if (m_source.is_globally_masked(i, j)) continue;
+
+            auto this_ptr = mutable_ptr(i, j);
+            int distance = reset ? PatchDistanceMetric::kDistanceScale : this_ptr[2];
+            if (distance < PatchDistanceMetric::kDistanceScale) {
+                continue;
+            }
+
+            int i_target = 0, j_target = 0;
+            for (int t = 0; t < max_retry; ++t) {
+                i_target = rand() % this_size.height;
+                j_target = rand() % this_size.width;
+                if (m_target.is_globally_masked(i_target, j_target)) continue;
+
+                distance = _distance(i, j, i_target, j_target);
+                if (distance < PatchDistanceMetric::kDistanceScale)
+                    break;
+            }
+
+            this_ptr[0] = i_target, this_ptr[1] = j_target, this_ptr[2] = distance;
+        }
+    }
+}
+
+void NearestNeighborField::_initialize_field_from(const NearestNeighborField &other, int max_retry) {
+    const auto &this_size = source_size();
+    const auto &other_size = other.source_size();
+    double fi = static_cast<double>(this_size.height) / other_size.height;
+    double fj = static_cast<double>(this_size.width) / other_size.width;
+
+    for (int i = 0; i < this_size.height; ++i) {
+        for (int j = 0; j < this_size.width; ++j) {
+            if (m_source.is_globally_masked(i, j)) continue;
+
+            int ilow = static_cast<int>(std::min(i / fi, static_cast<double>(other_size.height - 1)));
+            int jlow = static_cast<int>(std::min(j / fj, static_cast<double>(other_size.width - 1)));
+            auto this_value = mutable_ptr(i, j);
+            auto other_value = other.ptr(ilow, jlow);
+
+            this_value[0] = static_cast<int>(other_value[0] * fi);
+            this_value[1] = static_cast<int>(other_value[1] * fj);
+            this_value[2] = _distance(i, j, this_value[0], this_value[1]);
+        }
+    }
+
+    _randomize_field(max_retry, false);
+}
+
+void NearestNeighborField::minimize(int nr_pass) {
+    const auto &this_size = source_size();
+    while (nr_pass--) {
+        for (int i = 0; i < this_size.height; ++i)
+            for (int j = 0; j < this_size.width; ++j) {
+                if (m_source.is_globally_masked(i, j)) continue;
+                if (at(i, j, 2) > 0) _minimize_link(i, j, +1);
+            }
+        for (int i = this_size.height - 1; i >= 0; --i)
+            for (int j = this_size.width - 1; j >= 0; --j) {
+                if (m_source.is_globally_masked(i, j)) continue;
+                if (at(i, j, 2) > 0) _minimize_link(i, j, -1);
+            }
+    }
+}
+
+void NearestNeighborField::_minimize_link(int y, int x, int direction) {
+    const auto &this_size = source_size();
+    const auto &this_target_size = target_size();
+    auto this_ptr = mutable_ptr(y, x);
+
+    // propagation along the y direction.
+    if (y - direction >= 0 && y - direction < this_size.height && !m_source.is_globally_masked(y - direction, x)) {
+        int yp = at(y - direction, x, 0) + direction;
+        int xp = at(y - direction, x, 1);
+        int dp = _distance(y, x, yp, xp);
+        if (dp < at(y, x, 2)) {
+            this_ptr[0] = yp, this_ptr[1] = xp, this_ptr[2] = dp;
+        }
+    }
+
+    // propagation along the x direction.
+    if (x - direction >= 0 && x - direction < this_size.width && !m_source.is_globally_masked(y, x - direction)) {
+        int yp = at(y, x - direction, 0);
+        int xp = at(y, x - direction, 1) + direction;
+        int dp = _distance(y, x, yp, xp);
+        if (dp < at(y, x, 2)) {
+            this_ptr[0] = yp, this_ptr[1] = xp, this_ptr[2] = dp;
+        }
+    }
+
+    // random search with a progressive step size.
+    int random_scale = (std::min(this_target_size.height, this_target_size.width) - 1) / 2;
+    while (random_scale > 0) {
+        int yp = this_ptr[0] + (rand() % (2 * random_scale + 1) - random_scale);
+        int xp = this_ptr[1] + (rand() % (2 * random_scale + 1) - random_scale);
+        yp = clamp(yp, 0, target_size().height - 1);
+        xp = clamp(xp, 0, target_size().width - 1);
+
+        if (m_target.is_globally_masked(yp, xp)) {
+            random_scale /= 2;
+        }
+
+        int dp = _distance(y, x, yp, xp);
+        if (dp < at(y, x, 2)) {
+            this_ptr[0] = yp, this_ptr[1] = xp, this_ptr[2] = dp;
+        }
+        random_scale /= 2;
+    }
+}
+
+const int PatchDistanceMetric::kDistanceScale = 65535;
+const int PatchSSDDistanceMetric::kSSDScale = 9 * 255 * 255;
+
+namespace {
+
+inline int pow2(int i) {
+    return i * i;
+}
+
+int distance_masked_images(
+    const MaskedImage &source, int ys, int xs,
+    const MaskedImage &target, int yt, int xt,
+    int patch_size
+) {
+    long double distance = 0;
+    long double wsum = 0;
+
+    source.compute_image_gradients();
+    target.compute_image_gradients();
+
+    auto source_size = source.size();
+    auto target_size = target.size();
+
+    for (int dy = -patch_size; dy <= patch_size; ++dy) {
+        const int yys = ys + dy, yyt = yt + dy;
+
+        if (yys <= 0 || yys >= source_size.height - 1 || yyt <= 0 || yyt >= target_size.height - 1) {
+            distance += (long double)(PatchSSDDistanceMetric::kSSDScale) * (2 * patch_size + 1);
+            wsum += 2 * patch_size + 1;
+            continue;
+        }
+
+        const auto *p_si = source.image().ptr<unsigned char>(yys, 0);
+        const auto *p_ti = target.image().ptr<unsigned char>(yyt, 0);
+        const auto *p_sm = source.mask().ptr<unsigned char>(yys, 0);
+        const auto *p_tm = target.mask().ptr<unsigned char>(yyt, 0);
+
+        const unsigned char *p_sgm = nullptr;
+        const unsigned char *p_tgm = nullptr;
+        if (!source.global_mask().empty()) {
+            p_sgm = source.global_mask().ptr<unsigned char>(yys, 0);
+            p_tgm = target.global_mask().ptr<unsigned char>(yyt, 0);
+        }
+
+        const auto *p_sgy = source.grady().ptr<unsigned char>(yys, 0);
+        const auto *p_tgy = target.grady().ptr<unsigned char>(yyt, 0);
+        const auto *p_sgx = source.gradx().ptr<unsigned char>(yys, 0);
+        const auto *p_tgx = target.gradx().ptr<unsigned char>(yyt, 0);
+
+        for (int dx = -patch_size; dx <= patch_size; ++dx) {
+            int xxs = xs + dx, xxt = xt + dx;
+            wsum += 1;
+
+            if (xxs <= 0 || xxs >= source_size.width - 1 || xxt <= 0 || xxt >= source_size.width - 1) {
+                distance += PatchSSDDistanceMetric::kSSDScale;
+                continue;
+            }
+
+            if (p_sm[xxs] || p_tm[xxt] || (p_sgm && p_sgm[xxs]) || (p_tgm && p_tgm[xxt]) ) {
+                distance += PatchSSDDistanceMetric::kSSDScale;
+                continue;
+            }
+
+            int ssd = 0;
+            for (int c = 0; c < 3; ++c) {
+                int s_value = p_si[xxs * 3 + c];
+                int t_value = p_ti[xxt * 3 + c];
+                int s_gy = p_sgy[xxs * 3 + c];
+                int t_gy = p_tgy[xxt * 3 + c];
+                int s_gx = p_sgx[xxs * 3 + c];
+                int t_gx = p_tgx[xxt * 3 + c];
+
+                ssd += pow2(static_cast<int>(s_value) - t_value);
+                ssd += pow2(static_cast<int>(s_gx) - t_gx);
+                ssd += pow2(static_cast<int>(s_gy) - t_gy);
+            }
+            distance += ssd;
+        }
+    }
+
+    distance /= (long double)(PatchSSDDistanceMetric::kSSDScale);
+
+    int res = int(PatchDistanceMetric::kDistanceScale * distance / wsum);
+    if (res < 0 || res > PatchDistanceMetric::kDistanceScale) return PatchDistanceMetric::kDistanceScale;
+    return res;
+}
+
+}
+
+int PatchSSDDistanceMetric::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
+    return distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
+}
+
+int DebugPatchSSDDistanceMetric::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
+    fprintf(stderr, "DebugPatchSSDDistanceMetric: %d %d %d %d\n", source.size().width, source.size().height, m_width, m_height);
+    return distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
+}
+
+int RegularityGuidedPatchDistanceMetricV1::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
+    double dx = remainder(double(source_x - target_x) / source.size().width, m_dx1);
+    double dy = remainder(double(source_y - target_y) / source.size().height, m_dy2);
+
+    double score1 = sqrt(dx * dx + dy *dy) / m_scale;
+    if (score1 < 0 || score1 > 1) score1 = 1;
+    score1 *= PatchDistanceMetric::kDistanceScale;
+
+    double score2 = distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
+    double score = score1 * m_weight + score2 / (1 + m_weight);
+    return static_cast<int>(score / (1 + m_weight));
+}
+
+int RegularityGuidedPatchDistanceMetricV2::operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const {
+    if (target_y < 0 || target_y >= target.size().height || target_x < 0 || target_x >= target.size().width)
+        return PatchDistanceMetric::kDistanceScale;
+
+    int source_scale = m_ijmap.size().height / source.size().height;
+    int target_scale = m_ijmap.size().height / target.size().height;
+
+    // fprintf(stderr, "RegularityGuidedPatchDistanceMetricV2 %d %d %d %d\n", source_y * source_scale, m_ijmap.size().height, source_x * source_scale, m_ijmap.size().width);
+
+    double score1 = PatchDistanceMetric::kDistanceScale;
+    if (!source.is_globally_masked(source_y, source_x) && !target.is_globally_masked(target_y, target_x)) {
+        auto source_ij = m_ijmap.ptr<float>(source_y * source_scale, source_x * source_scale);
+        auto target_ij = m_ijmap.ptr<float>(target_y * target_scale, target_x * target_scale);
+
+        float di = fabs(source_ij[0] - target_ij[0]); if (di > 0.5) di = 1 - di;
+        float dj = fabs(source_ij[1] - target_ij[1]); if (dj > 0.5) dj = 1 - dj;
+        score1 = sqrt(di * di + dj *dj) / 0.707;
+        if (score1 < 0 || score1 > 1) score1 = 1;
+        score1 *= PatchDistanceMetric::kDistanceScale;
+    }
+
+    double score2 = distance_masked_images(source, source_y, source_x, target, target_y, target_x, m_patch_size);
+    double score = score1 * m_weight + score2;
+    return int(score / (1 + m_weight));
+}
+

PyPatchMatch/csrc/nnf.h

@@ -0,0 +1,133 @@
+#pragma once
+
+#include <opencv2/core.hpp>
+#include "masked_image.h"
+
+class PatchDistanceMetric {
+public:
+    PatchDistanceMetric(int patch_size) : m_patch_size(patch_size) {}
+    virtual ~PatchDistanceMetric() = default;
+
+    inline int patch_size() const { return m_patch_size; }
+    virtual int operator()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const = 0;
+    static const int kDistanceScale;
+
+protected:
+    int m_patch_size;
+};
+
+class NearestNeighborField {
+public:
+    NearestNeighborField() : m_source(), m_target(), m_field(), m_distance_metric(nullptr) {
+        // pass
+    }
+    NearestNeighborField(const MaskedImage &source, const MaskedImage &target, const PatchDistanceMetric *metric, int max_retry = 20)
+        : m_source(source), m_target(target), m_distance_metric(metric) {
+        m_field = cv::Mat(m_source.size(), CV_32SC3);
+        _randomize_field(max_retry);
+    }
+    NearestNeighborField(const MaskedImage &source, const MaskedImage &target, const PatchDistanceMetric *metric, const NearestNeighborField &other, int max_retry = 20)
+            : m_source(source), m_target(target), m_distance_metric(metric) {
+        m_field = cv::Mat(m_source.size(), CV_32SC3);
+        _initialize_field_from(other, max_retry);
+    }
+
+    const MaskedImage &source() const {
+        return m_source;
+    }
+    const MaskedImage &target() const {
+        return m_target;
+    }
+    inline cv::Size source_size() const {
+        return m_source.size();
+    }
+    inline cv::Size target_size() const {
+        return m_target.size();
+    }
+    inline void set_source(const MaskedImage &source) {
+        m_source = source;
+    }
+    inline void set_target(const MaskedImage &target) {
+        m_target = target;
+    }
+
+    inline int *mutable_ptr(int y, int x) {
+        return m_field.ptr<int>(y, x);
+    }
+    inline const int *ptr(int y, int x) const {
+        return m_field.ptr<int>(y, x);
+    }
+
+    inline int at(int y, int x, int c) const {
+        return m_field.ptr<int>(y, x)[c];
+    }
+    inline int &at(int y, int x, int c) {
+        return m_field.ptr<int>(y, x)[c];
+    }
+    inline void set_identity(int y, int x) {
+        auto ptr = mutable_ptr(y, x);
+        ptr[0] = y, ptr[1] = x, ptr[2] = 0;
+    }
+
+    void minimize(int nr_pass);
+
+private:
+    inline int _distance(int source_y, int source_x, int target_y, int target_x) {
+        return (*m_distance_metric)(m_source, source_y, source_x, m_target, target_y, target_x);
+    }
+
+    void _randomize_field(int max_retry = 20, bool reset = true);
+    void _initialize_field_from(const NearestNeighborField &other, int max_retry);
+    void _minimize_link(int y, int x, int direction);
+
+    MaskedImage m_source;
+    MaskedImage m_target;
+    cv::Mat m_field;  // { y_target, x_target, distance_scaled }
+    const PatchDistanceMetric *m_distance_metric;
+};
+
+
+class PatchSSDDistanceMetric : public PatchDistanceMetric {
+public:
+    using PatchDistanceMetric::PatchDistanceMetric;
+    virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
+    static const int kSSDScale;
+};
+
+class DebugPatchSSDDistanceMetric : public PatchDistanceMetric {
+public:
+    DebugPatchSSDDistanceMetric(int patch_size, int width, int height) : PatchDistanceMetric(patch_size), m_width(width), m_height(height) {}
+    virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
+protected:
+    int m_width, m_height;
+};
+
+class RegularityGuidedPatchDistanceMetricV1 : public PatchDistanceMetric {
+public:
+    RegularityGuidedPatchDistanceMetricV1(int patch_size, double dx1, double dy1, double dx2, double dy2, double weight)
+        : PatchDistanceMetric(patch_size), m_dx1(dx1), m_dy1(dy1), m_dx2(dx2), m_dy2(dy2), m_weight(weight) {
+
+        assert(m_dy1 == 0);
+        assert(m_dx2 == 0);
+        m_scale = sqrt(m_dx1 * m_dx1 + m_dy2 * m_dy2) / 4;
+    }
+    virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
+
+protected:
+    double m_dx1, m_dy1, m_dx2, m_dy2;
+    double m_scale, m_weight;
+};
+
+class RegularityGuidedPatchDistanceMetricV2 : public PatchDistanceMetric {
+public:
+    RegularityGuidedPatchDistanceMetricV2(int patch_size, cv::Mat ijmap, double weight)
+        : PatchDistanceMetric(patch_size), m_ijmap(ijmap), m_weight(weight) {
+
+    }
+    virtual int operator ()(const MaskedImage &source, int source_y, int source_x, const MaskedImage &target, int target_y, int target_x) const;
+
+protected:
+    cv::Mat m_ijmap;
+    double m_width, m_height, m_weight;
+};
+

PyPatchMatch/csrc/pyinterface.cpp

@@ -0,0 +1,107 @@
+#include "pyinterface.h"
+#include "inpaint.h"
+
+static unsigned int PM_seed = 1212;
+static bool PM_verbose = false;
+
+int _dtype_py_to_cv(int dtype_py);
+int _dtype_cv_to_py(int dtype_cv);
+cv::Mat _py_to_cv2(PM_mat_t pymat);
+PM_mat_t _cv2_to_py(cv::Mat cvmat);
+
+void PM_set_random_seed(unsigned int seed) {
+    PM_seed = seed;
+}
+
+void PM_set_verbose(int value) {
+    PM_verbose = static_cast<bool>(value);
+}
+
+void PM_free_pymat(PM_mat_t pymat) {
+    free(pymat.data_ptr);
+}
+
+PM_mat_t PM_inpaint(PM_mat_t source_py, PM_mat_t mask_py, int patch_size) {
+    cv::Mat source = _py_to_cv2(source_py);
+    cv::Mat mask = _py_to_cv2(mask_py);
+    auto metric = PatchSSDDistanceMetric(patch_size);
+    cv::Mat result = Inpainting(source, mask, &metric).run(PM_verbose, false, PM_seed);
+    return _cv2_to_py(result);
+}
+
+PM_mat_t PM_inpaint_regularity(PM_mat_t source_py, PM_mat_t mask_py, PM_mat_t ijmap_py, int patch_size, float guide_weight) {
+    cv::Mat source = _py_to_cv2(source_py);
+    cv::Mat mask = _py_to_cv2(mask_py);
+    cv::Mat ijmap = _py_to_cv2(ijmap_py);
+
+    auto metric = RegularityGuidedPatchDistanceMetricV2(patch_size, ijmap, guide_weight);
+    cv::Mat result = Inpainting(source, mask, &metric).run(PM_verbose, false, PM_seed);
+    return _cv2_to_py(result);
+}
+
+PM_mat_t PM_inpaint2(PM_mat_t source_py, PM_mat_t mask_py, PM_mat_t global_mask_py, int patch_size) {
+    cv::Mat source = _py_to_cv2(source_py);
+    cv::Mat mask = _py_to_cv2(mask_py);
+    cv::Mat global_mask = _py_to_cv2(global_mask_py);
+
+    auto metric = PatchSSDDistanceMetric(patch_size);
+    cv::Mat result = Inpainting(source, mask, global_mask, &metric).run(PM_verbose, false, PM_seed);
+    return _cv2_to_py(result);
+}
+
+PM_mat_t PM_inpaint2_regularity(PM_mat_t source_py, PM_mat_t mask_py, PM_mat_t global_mask_py, PM_mat_t ijmap_py, int patch_size, float guide_weight) {
+    cv::Mat source = _py_to_cv2(source_py);
+    cv::Mat mask = _py_to_cv2(mask_py);
+    cv::Mat global_mask = _py_to_cv2(global_mask_py);
+    cv::Mat ijmap = _py_to_cv2(ijmap_py);
+
+    auto metric = RegularityGuidedPatchDistanceMetricV2(patch_size, ijmap, guide_weight);
+    cv::Mat result = Inpainting(source, mask, global_mask, &metric).run(PM_verbose, false, PM_seed);
+    return _cv2_to_py(result);
+}
+
+int _dtype_py_to_cv(int dtype_py) {
+    switch (dtype_py) {
+        case PM_UINT8: return CV_8U;
+        case PM_INT8: return CV_8S;
+        case PM_UINT16: return CV_16U;
+        case PM_INT16: return CV_16S;
+        case PM_INT32: return CV_32S;
+        case PM_FLOAT32: return CV_32F;
+        case PM_FLOAT64: return CV_64F;
+    }
+
+    return CV_8U;
+}
+
+int _dtype_cv_to_py(int dtype_cv) {
+    switch (dtype_cv) {
+        case CV_8U: return PM_UINT8;
+        case CV_8S: return PM_INT8;
+        case CV_16U: return PM_UINT16;
+        case CV_16S: return PM_INT16;
+        case CV_32S: return PM_INT32;
+        case CV_32F: return PM_FLOAT32;
+        case CV_64F: return PM_FLOAT64;
+    }
+
+    return PM_UINT8;
+}
+
+cv::Mat _py_to_cv2(PM_mat_t pymat) {
+    int dtype = _dtype_py_to_cv(pymat.dtype);
+    dtype = CV_MAKETYPE(pymat.dtype, pymat.shape.channels);
+    return cv::Mat(cv::Size(pymat.shape.width, pymat.shape.height), dtype, pymat.data_ptr).clone();
+}
+
+PM_mat_t _cv2_to_py(cv::Mat cvmat) {
+    PM_shape_t shape = {cvmat.size().width, cvmat.size().height, cvmat.channels()};
+    int dtype = _dtype_cv_to_py(cvmat.depth());
+    size_t dsize = cvmat.total() * cvmat.elemSize();
+
+    void *data_ptr = reinterpret_cast<void *>(malloc(dsize));
+    memcpy(data_ptr, reinterpret_cast<void *>(cvmat.data), dsize);
+
+    return PM_mat_t {data_ptr, shape, dtype};
+}
+

PyPatchMatch/csrc/pyinterface.h

@@ -0,0 +1,38 @@
+#include <opencv2/core.hpp>
+#include <cstdlib>
+#include <cstdio>
+#include <cstring>
+
+extern "C" {
+
+struct PM_shape_t {
+    int width, height, channels;
+};
+
+enum PM_dtype_e {
+    PM_UINT8,
+    PM_INT8,
+    PM_UINT16,
+    PM_INT16,
+    PM_INT32,
+    PM_FLOAT32,
+    PM_FLOAT64,
+};
+
+struct PM_mat_t {
+    void *data_ptr;
+    PM_shape_t shape;
+    int dtype;
+};
+
+void PM_set_random_seed(unsigned int seed);
+void PM_set_verbose(int value);
+
+void PM_free_pymat(PM_mat_t pymat);
+PM_mat_t PM_inpaint(PM_mat_t image, PM_mat_t mask, int patch_size);
+PM_mat_t PM_inpaint_regularity(PM_mat_t image, PM_mat_t mask, PM_mat_t ijmap, int patch_size, float guide_weight);
+PM_mat_t PM_inpaint2(PM_mat_t image, PM_mat_t mask, PM_mat_t global_mask, int patch_size);
+PM_mat_t PM_inpaint2_regularity(PM_mat_t image, PM_mat_t mask, PM_mat_t global_mask, PM_mat_t ijmap, int patch_size, float guide_weight);
+
+} /*  extern "C" */
+

PyPatchMatch/examples/.gitignore

@@ -0,0 +1,2 @@
+/cpp_example.exe
+/images/*recovered.bmp

PyPatchMatch/examples/cpp_example.cpp

@@ -0,0 +1,31 @@
+#include <iostream>
+#include <opencv2/imgcodecs.hpp>
+#include <opencv2/highgui.hpp>
+
+#include "masked_image.h"
+#include "nnf.h"
+#include "inpaint.h"
+
+int main() {
+    auto source = cv::imread("./images/forest_pruned.bmp", cv::IMREAD_COLOR);
+
+    auto mask = cv::Mat(source.size(), CV_8UC1);
+    mask = cv::Scalar::all(0);
+    for (int i = 0; i < source.size().height; ++i) {
+        for (int j = 0; j < source.size().width; ++j) {
+            auto source_ptr = source.ptr<unsigned char>(i, j);
+            if (source_ptr[0] == 255 && source_ptr[1] == 255 && source_ptr[2] == 255) {
+                mask.at<unsigned char>(i, j) = 1;
+            }
+        }
+    }
+
+    auto metric = PatchSSDDistanceMetric(3);
+    auto result = Inpainting(source, mask, &metric).run(true, true);
+    // cv::imwrite("./images/forest_recovered.bmp", result);
+    // cv::imshow("Result", result);
+    // cv::waitKey();
+
+    return 0;
+}
+

PyPatchMatch/examples/cpp_example_run.sh

@@ -0,0 +1,18 @@
+#! /bin/bash
+#
+# cpp_example_run.sh
+# Copyright (C) 2020 Jiayuan Mao <maojiayuan@gmail.com>
+#
+# Distributed under terms of the MIT license.
+#
+
+set -x
+
+CFLAGS="-std=c++14 -O2 $(pkg-config --cflags opencv)"
+LDFLAGS="$(pkg-config --libs opencv)"
+g++ $CFLAGS cpp_example.cpp -I../csrc/ -L../ -lpatchmatch $LDFLAGS -o cpp_example.exe
+
+export DYLD_LIBRARY_PATH=../:$DYLD_LIBRARY_PATH  # For macOS
+export LD_LIBRARY_PATH=../:$LD_LIBRARY_PATH  # For Linux
+time ./cpp_example.exe
+

PyPatchMatch/examples/py_example.py

@@ -0,0 +1,21 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+# File   : test.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 01/09/2020
+#
+# Distributed under terms of the MIT license.
+
+from PIL import Image
+
+import sys
+sys.path.insert(0, '../')
+import patch_match
+
+
+if __name__ == '__main__':
+    source = Image.open('./images/forest_pruned.bmp')
+    result = patch_match.inpaint(source, patch_size=3)
+    Image.fromarray(result).save('./images/forest_recovered.bmp')
+

PyPatchMatch/examples/py_example_global_mask.py

@@ -0,0 +1,27 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+# File   : test.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 01/09/2020
+#
+# Distributed under terms of the MIT license.
+
+import numpy as np
+from PIL import Image
+
+import sys
+sys.path.insert(0, '../')
+import patch_match
+
+
+if __name__ == '__main__':
+    patch_match.set_verbose(True)
+    source = Image.open('./images/forest_pruned.bmp')
+    source = np.array(source)
+    source[:100, :100] = 255
+    global_mask = np.zeros_like(source[..., 0])
+    global_mask[:100, :100] = 1
+    result = patch_match.inpaint(source, global_mask=global_mask, patch_size=3)
+    Image.fromarray(result).save('./images/forest_recovered.bmp')
+

PyPatchMatch/patch_match.py

@@ -0,0 +1,263 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+# File   : patch_match.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 01/09/2020
+#
+# Distributed under terms of the MIT license.
+
+import ctypes
+import os.path as osp
+from typing import Optional, Union
+
+import numpy as np
+from PIL import Image
+
+
+import os
+if os.name!="nt":
+    # Otherwise, fall back to the subprocess.
+    import subprocess
+    print('Compiling and loading c extensions from "{}".'.format(osp.realpath(osp.dirname(__file__))))
+    # subprocess.check_call(['./travis.sh'], cwd=osp.dirname(__file__))
+    subprocess.check_call("make clean && make", cwd=osp.dirname(__file__), shell=True)
+
+
+__all__ = ['set_random_seed', 'set_verbose', 'inpaint', 'inpaint_regularity']
+
+
+class CShapeT(ctypes.Structure):
+    _fields_ = [
+        ('width', ctypes.c_int),
+        ('height', ctypes.c_int),
+        ('channels', ctypes.c_int),
+    ]
+
+
+class CMatT(ctypes.Structure):
+    _fields_ = [
+        ('data_ptr', ctypes.c_void_p),
+        ('shape', CShapeT),
+        ('dtype', ctypes.c_int)
+    ]
+
+import tempfile
+from urllib.request import urlopen, Request
+import shutil
+from pathlib import Path
+from tqdm import tqdm
+
+def download_url_to_file(url, dst, hash_prefix=None, progress=True):
+    r"""Download object at the given URL to a local path.
+
+    Args:
+        url (string): URL of the object to download
+        dst (string): Full path where object will be saved, e.g. ``/tmp/temporary_file``
+        hash_prefix (string, optional): If not None, the SHA256 downloaded file should start with ``hash_prefix``.
+            Default: None
+        progress (bool, optional): whether or not to display a progress bar to stderr
+            Default: True
+    https://pytorch.org/docs/stable/_modules/torch/hub.html#load_state_dict_from_url
+    """
+    file_size = None
+    req = Request(url)
+    u = urlopen(req)
+    meta = u.info()
+    if hasattr(meta, 'getheaders'):
+        content_length = meta.getheaders("Content-Length")
+    else:
+        content_length = meta.get_all("Content-Length")
+    if content_length is not None and len(content_length) > 0:
+        file_size = int(content_length[0])
+
+    # We deliberately save it in a temp file and move it after
+    # download is complete. This prevents a local working checkpoint
+    # being overridden by a broken download.
+    dst = os.path.expanduser(dst)
+    dst_dir = os.path.dirname(dst)
+    f = tempfile.NamedTemporaryFile(delete=False, dir=dst_dir)
+
+    try:
+        with tqdm(total=file_size, disable=not progress,
+                  unit='B', unit_scale=True, unit_divisor=1024) as pbar:
+            while True:
+                buffer = u.read(8192)
+                if len(buffer) == 0:
+                    break
+                f.write(buffer)
+                pbar.update(len(buffer))
+
+        f.close()
+        shutil.move(f.name, dst)
+    finally:
+        f.close()
+        if os.path.exists(f.name):
+            os.remove(f.name)
+
+if os.name!="nt":
+    PMLIB = ctypes.CDLL(osp.join(osp.dirname(__file__), 'libpatchmatch.so'))
+else:
+    if not os.path.exists(osp.join(osp.dirname(__file__), 'libpatchmatch.dll')):
+        download_url_to_file(url="https://github.com/lkwq007/PyPatchMatch/releases/download/v0.1/libpatchmatch.dll",dst=osp.join(osp.dirname(__file__), 'libpatchmatch.dll'))
+    if not os.path.exists(osp.join(osp.dirname(__file__), 'opencv_world460.dll')):
+        download_url_to_file(url="https://github.com/lkwq007/PyPatchMatch/releases/download/v0.1/opencv_world460.dll",dst=osp.join(osp.dirname(__file__), 'opencv_world460.dll'))
+    if not os.path.exists(osp.join(osp.dirname(__file__), 'libpatchmatch.dll')):
+        print("[Dependency Missing] Please download https://github.com/lkwq007/PyPatchMatch/releases/download/v0.1/libpatchmatch.dll and put it into the PyPatchMatch folder")
+    if not os.path.exists(osp.join(osp.dirname(__file__), 'opencv_world460.dll')):
+        print("[Dependency Missing] Please download https://github.com/lkwq007/PyPatchMatch/releases/download/v0.1/opencv_world460.dll and put it into the PyPatchMatch folder")
+    PMLIB = ctypes.CDLL(osp.join(osp.dirname(__file__), 'libpatchmatch.dll'))
+
+PMLIB.PM_set_random_seed.argtypes = [ctypes.c_uint]
+PMLIB.PM_set_verbose.argtypes = [ctypes.c_int]
+PMLIB.PM_free_pymat.argtypes = [CMatT]
+PMLIB.PM_inpaint.argtypes = [CMatT, CMatT, ctypes.c_int]
+PMLIB.PM_inpaint.restype = CMatT
+PMLIB.PM_inpaint_regularity.argtypes = [CMatT, CMatT, CMatT, ctypes.c_int, ctypes.c_float]
+PMLIB.PM_inpaint_regularity.restype = CMatT
+PMLIB.PM_inpaint2.argtypes = [CMatT, CMatT, CMatT, ctypes.c_int]
+PMLIB.PM_inpaint2.restype = CMatT
+PMLIB.PM_inpaint2_regularity.argtypes = [CMatT, CMatT, CMatT, CMatT, ctypes.c_int, ctypes.c_float]
+PMLIB.PM_inpaint2_regularity.restype = CMatT
+
+
+def set_random_seed(seed: int):
+    PMLIB.PM_set_random_seed(ctypes.c_uint(seed))
+
+
+def set_verbose(verbose: bool):
+    PMLIB.PM_set_verbose(ctypes.c_int(verbose))
+
+
+def inpaint(
+    image: Union[np.ndarray, Image.Image],
+    mask: Optional[Union[np.ndarray, Image.Image]] = None,
+    *,
+    global_mask: Optional[Union[np.ndarray, Image.Image]] = None,
+    patch_size: int = 15
+) -> np.ndarray:
+    """
+    PatchMatch based inpainting proposed in:
+
+        PatchMatch : A Randomized Correspondence Algorithm for Structural Image Editing
+        C.Barnes, E.Shechtman, A.Finkelstein and Dan B.Goldman
+        SIGGRAPH 2009
+
+    Args:
+        image (Union[np.ndarray, Image.Image]): the input image, should be 3-channel RGB/BGR.
+        mask (Union[np.array, Image.Image], optional): the mask of the hole(s) to be filled, should be 1-channel.
+        If not provided (None), the algorithm will treat all purely white pixels as the holes (255, 255, 255).
+        global_mask (Union[np.array, Image.Image], optional): the target mask of the output image.
+        patch_size (int): the patch size for the inpainting algorithm.
+
+    Return:
+        result (np.ndarray): the repaired image, of the same size as the input image.
+    """
+
+    if isinstance(image, Image.Image):
+        image = np.array(image)
+    image = np.ascontiguousarray(image)
+    assert image.ndim == 3 and image.shape[2] == 3 and image.dtype == 'uint8'
+
+    if mask is None:
+        mask = (image == (255, 255, 255)).all(axis=2, keepdims=True).astype('uint8')
+        mask = np.ascontiguousarray(mask)
+    else:
+        mask = _canonize_mask_array(mask)
+
+    if global_mask is None:
+        ret_pymat = PMLIB.PM_inpaint(np_to_pymat(image), np_to_pymat(mask), ctypes.c_int(patch_size))
+    else:
+        global_mask = _canonize_mask_array(global_mask)
+        ret_pymat = PMLIB.PM_inpaint2(np_to_pymat(image), np_to_pymat(mask), np_to_pymat(global_mask), ctypes.c_int(patch_size))
+
+    ret_npmat = pymat_to_np(ret_pymat)
+    PMLIB.PM_free_pymat(ret_pymat)
+
+    return ret_npmat
+
+
+def inpaint_regularity(
+    image: Union[np.ndarray, Image.Image],
+    mask: Optional[Union[np.ndarray, Image.Image]],
+    ijmap: np.ndarray,
+    *,
+    global_mask: Optional[Union[np.ndarray, Image.Image]] = None,
+    patch_size: int = 15, guide_weight: float = 0.25
+) -> np.ndarray:
+    if isinstance(image, Image.Image):
+        image = np.array(image)
+    image = np.ascontiguousarray(image)
+
+    assert isinstance(ijmap, np.ndarray) and ijmap.ndim == 3 and ijmap.shape[2] == 3 and ijmap.dtype == 'float32'
+    ijmap = np.ascontiguousarray(ijmap)
+
+    assert image.ndim == 3 and image.shape[2] == 3 and image.dtype == 'uint8'
+    if mask is None:
+        mask = (image == (255, 255, 255)).all(axis=2, keepdims=True).astype('uint8')
+        mask = np.ascontiguousarray(mask)
+    else:
+        mask = _canonize_mask_array(mask)
+
+
+    if global_mask is None:
+        ret_pymat = PMLIB.PM_inpaint_regularity(np_to_pymat(image), np_to_pymat(mask), np_to_pymat(ijmap), ctypes.c_int(patch_size), ctypes.c_float(guide_weight))
+    else:
+        global_mask = _canonize_mask_array(global_mask)
+        ret_pymat = PMLIB.PM_inpaint2_regularity(np_to_pymat(image), np_to_pymat(mask), np_to_pymat(global_mask), np_to_pymat(ijmap), ctypes.c_int(patch_size), ctypes.c_float(guide_weight))
+
+    ret_npmat = pymat_to_np(ret_pymat)
+    PMLIB.PM_free_pymat(ret_pymat)
+
+    return ret_npmat
+
+
+def _canonize_mask_array(mask):
+    if isinstance(mask, Image.Image):
+        mask = np.array(mask)
+    if mask.ndim == 2 and mask.dtype == 'uint8':
+        mask = mask[..., np.newaxis]
+    assert mask.ndim == 3 and mask.shape[2] == 1 and mask.dtype == 'uint8'
+    return np.ascontiguousarray(mask)
+
+
+dtype_pymat_to_ctypes = [
+    ctypes.c_uint8,
+    ctypes.c_int8,
+    ctypes.c_uint16,
+    ctypes.c_int16,
+    ctypes.c_int32,
+    ctypes.c_float,
+    ctypes.c_double,
+]
+
+
+dtype_np_to_pymat = {
+    'uint8': 0,
+    'int8': 1,
+    'uint16': 2,
+    'int16': 3,
+    'int32': 4,
+    'float32': 5,
+    'float64': 6,
+}
+
+
+def np_to_pymat(npmat):
+    assert npmat.ndim == 3
+    return CMatT(
+        ctypes.cast(npmat.ctypes.data, ctypes.c_void_p),
+        CShapeT(npmat.shape[1], npmat.shape[0], npmat.shape[2]),
+        dtype_np_to_pymat[str(npmat.dtype)]
+    )
+
+
+def pymat_to_np(pymat):
+    npmat = np.ctypeslib.as_array(
+        ctypes.cast(pymat.data_ptr, ctypes.POINTER(dtype_pymat_to_ctypes[pymat.dtype])),
+        (pymat.shape.height, pymat.shape.width, pymat.shape.channels)
+    )
+    ret = np.empty(npmat.shape, npmat.dtype)
+    ret[:] = npmat
+    return ret
+

PyPatchMatch/travis.sh

@@ -0,0 +1,9 @@
+#! /bin/bash
+#
+# travis.sh
+# Copyright (C) 2020 Jiayuan Mao <maojiayuan@gmail.com>
+#
+# Distributed under terms of the MIT license.
+#
+
+make clean && make

README.md

@@ -0,0 +1,14 @@
+---
+title: Stablediffusion Infinity
+emoji: ♾️
+colorFrom: blue
+colorTo: purple
+sdk: gradio
+sdk_version: 3.10.1
+app_file: app.py
+pinned: true
+license: apache-2.0
+duplicated_from: lnyan/stablediffusion-infinity
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,1057 @@
+import subprocess
+# import os.path as osp
+import pip
+# pip.main(["install","-v","-U","git+https://github.com/facebookresearch/xformers.git@main#egg=xformers"])
+# subprocess.check_call("pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers", cwd=osp.dirname(__file__), shell=True)
+
+import io
+import base64
+import os
+import sys
+
+import numpy as np
+import torch
+from torch import autocast
+import diffusers
+from diffusers.configuration_utils import FrozenDict
+from diffusers import (
+    StableDiffusionPipeline,
+    StableDiffusionInpaintPipeline,
+    StableDiffusionImg2ImgPipeline,
+    StableDiffusionInpaintPipelineLegacy,
+    DDIMScheduler,
+    LMSDiscreteScheduler,
+    StableDiffusionUpscalePipeline,
+    DPMSolverMultistepScheduler
+)
+from diffusers.models import AutoencoderKL
+from PIL import Image
+from PIL import ImageOps
+import gradio as gr
+import base64
+import skimage
+import skimage.measure
+import yaml
+import json
+from enum import Enum
+
+try:
+    abspath = os.path.abspath(__file__)
+    dirname = os.path.dirname(abspath)
+    os.chdir(dirname)
+except:
+    pass
+
+from utils import *
+
+assert diffusers.__version__ >= "0.6.0", "Please upgrade diffusers to 0.6.0"
+
+USE_NEW_DIFFUSERS = True
+RUN_IN_SPACE = "RUN_IN_HG_SPACE" in os.environ
+
+
+class ModelChoice(Enum):
+    INPAINTING = "stablediffusion-inpainting"
+    INPAINTING_IMG2IMG = "stablediffusion-inpainting+img2img-v1.5"
+    MODEL_1_5 = "stablediffusion-v1.5"
+    MODEL_1_4 = "stablediffusion-v1.4"
+
+
+try:
+    from sd_grpcserver.pipeline.unified_pipeline import UnifiedPipeline
+except:
+    UnifiedPipeline = StableDiffusionInpaintPipeline
+
+# sys.path.append("./glid_3_xl_stable")
+
+USE_GLID = False
+# try:
+#     from glid3xlmodel import GlidModel
+# except:
+#     USE_GLID = False
+
+try:
+    cuda_available = torch.cuda.is_available()
+except:
+    cuda_available = False
+finally:
+    if sys.platform == "darwin":
+        device = "mps" if torch.backends.mps.is_available() else "cpu"
+    elif cuda_available:
+        device = "cuda"
+    else:
+        device = "cpu"
+
+import contextlib
+
+autocast = contextlib.nullcontext
+
+with open("config.yaml", "r") as yaml_in:
+    yaml_object = yaml.safe_load(yaml_in)
+    config_json = json.dumps(yaml_object)
+
+
+def load_html():
+    body, canvaspy = "", ""
+    with open("index.html", encoding="utf8") as f:
+        body = f.read()
+    with open("canvas.py", encoding="utf8") as f:
+        canvaspy = f.read()
+    body = body.replace("- paths:\n", "")
+    body = body.replace("  - ./canvas.py\n", "")
+    body = body.replace("from canvas import InfCanvas", canvaspy)
+    return body
+
+
+def test(x):
+    x = load_html()
+    return f"""<iframe id="sdinfframe" style="width: 100%; height: 600px" name="result" allow="midi; geolocation; microphone; camera; 
+    display-capture; encrypted-media; vertical-scroll 'none'" sandbox="allow-modals allow-forms 
+    allow-scripts allow-same-origin allow-popups 
+    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
+    allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
+
+
+DEBUG_MODE = False
+
+try:
+    SAMPLING_MODE = Image.Resampling.LANCZOS
+except Exception as e:
+    SAMPLING_MODE = Image.LANCZOS
+
+try:
+    contain_func = ImageOps.contain
+except Exception as e:
+
+    def contain_func(image, size, method=SAMPLING_MODE):
+        # from PIL: https://pillow.readthedocs.io/en/stable/reference/ImageOps.html#PIL.ImageOps.contain
+        im_ratio = image.width / image.height
+        dest_ratio = size[0] / size[1]
+        if im_ratio != dest_ratio:
+            if im_ratio > dest_ratio:
+                new_height = int(image.height / image.width * size[0])
+                if new_height != size[1]:
+                    size = (size[0], new_height)
+            else:
+                new_width = int(image.width / image.height * size[1])
+                if new_width != size[0]:
+                    size = (new_width, size[1])
+        return image.resize(size, resample=method)
+
+
+import argparse
+
+parser = argparse.ArgumentParser(description="stablediffusion-infinity")
+parser.add_argument("--port", type=int, help="listen port", dest="server_port")
+parser.add_argument("--host", type=str, help="host", dest="server_name")
+parser.add_argument("--share", action="store_true", help="share this app?")
+parser.add_argument("--debug", action="store_true", help="debug mode")
+parser.add_argument("--fp32", action="store_true", help="using full precision")
+parser.add_argument("--encrypt", action="store_true", help="using https?")
+parser.add_argument("--ssl_keyfile", type=str, help="path to ssl_keyfile")
+parser.add_argument("--ssl_certfile", type=str, help="path to ssl_certfile")
+parser.add_argument("--ssl_keyfile_password", type=str, help="ssl_keyfile_password")
+parser.add_argument(
+    "--auth", nargs=2, metavar=("username", "password"), help="use username password"
+)
+parser.add_argument(
+    "--remote_model",
+    type=str,
+    help="use a model (e.g. dreambooth fined) from huggingface hub",
+    default="",
+)
+parser.add_argument(
+    "--local_model", type=str, help="use a model stored on your PC", default=""
+)
+
+if __name__ == "__main__" and not RUN_IN_SPACE:
+    args = parser.parse_args()
+else:
+    args = parser.parse_args()
+# args = parser.parse_args(["--debug"])
+if args.auth is not None:
+    args.auth = tuple(args.auth)
+
+model = {}
+
+
+def get_token():
+    token = ""
+    if os.path.exists(".token"):
+        with open(".token", "r") as f:
+            token = f.read()
+    token = os.environ.get("hftoken", token)
+    return token
+
+
+def save_token(token):
+    with open(".token", "w") as f:
+        f.write(token)
+
+
+def prepare_scheduler(scheduler):
+    if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
+        new_config = dict(scheduler.config)
+        new_config["steps_offset"] = 1
+        scheduler._internal_dict = FrozenDict(new_config)
+    return scheduler
+
+
+def my_resize(width, height):
+    if width >= 512 and height >= 512:
+        return width, height
+    if width == height:
+        return 512, 512
+    smaller = min(width, height)
+    larger = max(width, height)
+    if larger >= 608:
+        return width, height
+    factor = 1
+    if smaller < 290:
+        factor = 2
+    elif smaller < 330:
+        factor = 1.75
+    elif smaller < 384:
+        factor = 1.375
+    elif smaller < 400:
+        factor = 1.25
+    elif smaller < 450:
+        factor = 1.125
+    return int(factor * width)//8*8, int(factor * height)//8*8
+
+
+def load_learned_embed_in_clip(
+    learned_embeds_path, text_encoder, tokenizer, token=None
+):
+    # https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/stable_conceptualizer_inference.ipynb
+    loaded_learned_embeds = torch.load(learned_embeds_path, map_location="cpu")
+
+    # separate token and the embeds
+    trained_token = list(loaded_learned_embeds.keys())[0]
+    embeds = loaded_learned_embeds[trained_token]
+
+    # cast to dtype of text_encoder
+    dtype = text_encoder.get_input_embeddings().weight.dtype
+    embeds.to(dtype)
+
+    # add the token in tokenizer
+    token = token if token is not None else trained_token
+    num_added_tokens = tokenizer.add_tokens(token)
+    if num_added_tokens == 0:
+        raise ValueError(
+            f"The tokenizer already contains the token {token}. Please pass a different `token` that is not already in the tokenizer."
+        )
+
+    # resize the token embeddings
+    text_encoder.resize_token_embeddings(len(tokenizer))
+
+    # get the id for the token and assign the embeds
+    token_id = tokenizer.convert_tokens_to_ids(token)
+    text_encoder.get_input_embeddings().weight.data[token_id] = embeds
+
+
+scheduler_dict = {"PLMS": None, "DDIM": None, "K-LMS": None, "DPM": None}
+
+
+class StableDiffusionInpaint:
+    def __init__(
+        self, token: str = "", model_name: str = "", model_path: str = "", **kwargs,
+    ):
+        self.token = token
+        original_checkpoint = False
+        if model_path and os.path.exists(model_path):
+            if model_path.endswith(".ckpt"):
+                original_checkpoint = True
+            elif model_path.endswith(".json"):
+                model_name = os.path.dirname(model_path)
+            else:
+                model_name = model_path
+        vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
+        vae.to(torch.float16)
+        if original_checkpoint:
+            print(f"Converting & Loading {model_path}")
+            from convert_checkpoint import convert_checkpoint
+
+            pipe = convert_checkpoint(model_path, inpainting=True)
+            if device == "cuda":
+                pipe.to(torch.float16)
+            inpaint = StableDiffusionInpaintPipeline(
+                vae=vae,
+                text_encoder=pipe.text_encoder,
+                tokenizer=pipe.tokenizer,
+                unet=pipe.unet,
+                scheduler=pipe.scheduler,
+                safety_checker=pipe.safety_checker,
+                feature_extractor=pipe.feature_extractor,
+            )
+        else:
+            print(f"Loading {model_name}")
+            if device == "cuda":
+                inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+                    model_name,
+                    revision="fp16",
+                    torch_dtype=torch.float16,
+                    use_auth_token=token,
+                    vae=vae
+                )
+            else:
+                inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+                    model_name, use_auth_token=token,
+                )
+        if os.path.exists("./embeddings"):
+            print("Note that StableDiffusionInpaintPipeline + embeddings is untested")
+            for item in os.listdir("./embeddings"):
+                if item.endswith(".bin"):
+                    load_learned_embed_in_clip(
+                        os.path.join("./embeddings", item),
+                        inpaint.text_encoder,
+                        inpaint.tokenizer,
+                    )
+        inpaint.to(device)
+        # try:
+            # inpaint.vae=torch.compile(inpaint.vae, dynamic=True)
+            # inpaint.unet=torch.compile(inpaint.unet, dynamic=True)
+        # except Exception as e:
+            # print(e)
+        # inpaint.enable_xformers_memory_efficient_attention()
+        # if device == "mps":
+        # _ = text2img("", num_inference_steps=1)
+        scheduler_dict["PLMS"] = inpaint.scheduler
+        scheduler_dict["DDIM"] = prepare_scheduler(
+            DDIMScheduler(
+                beta_start=0.00085,
+                beta_end=0.012,
+                beta_schedule="scaled_linear",
+                clip_sample=False,
+                set_alpha_to_one=False,
+            )
+        )
+        scheduler_dict["K-LMS"] = prepare_scheduler(
+            LMSDiscreteScheduler(
+                beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
+            )
+        )
+        scheduler_dict["DPM"] = prepare_scheduler(
+            DPMSolverMultistepScheduler.from_config(inpaint.scheduler.config)
+        )
+        self.safety_checker = inpaint.safety_checker
+        save_token(token)
+        try:
+            total_memory = torch.cuda.get_device_properties(0).total_memory // (
+                1024 ** 3
+            )
+            if total_memory <= 5:
+                inpaint.enable_attention_slicing()
+        except:
+            pass
+        self.inpaint = inpaint
+
+    def run(
+        self,
+        image_pil,
+        prompt="",
+        negative_prompt="",
+        guidance_scale=7.5,
+        resize_check=True,
+        enable_safety=True,
+        fill_mode="patchmatch",
+        strength=0.75,
+        step=50,
+        enable_img2img=False,
+        use_seed=False,
+        seed_val=-1,
+        generate_num=1,
+        scheduler="",
+        scheduler_eta=0.0,
+        **kwargs,
+    ):
+        inpaint = self.inpaint
+        selected_scheduler = scheduler_dict.get(scheduler, scheduler_dict["PLMS"])
+        for item in [inpaint]:
+            item.scheduler = selected_scheduler
+            if enable_safety:
+                item.safety_checker = self.safety_checker
+            else:
+                item.safety_checker = lambda images, **kwargs: (images, False)
+        width, height = image_pil.size
+        sel_buffer = np.array(image_pil)
+        img = sel_buffer[:, :, 0:3]
+        mask = sel_buffer[:, :, -1]
+        nmask = 255 - mask
+        process_width = width
+        process_height = height
+        if resize_check:
+            process_width, process_height = my_resize(width, height)
+        process_width=process_width*8//8
+        process_height=process_height*8//8
+        extra_kwargs = {
+            "num_inference_steps": step,
+            "guidance_scale": guidance_scale,
+            "eta": scheduler_eta,
+        }
+        if USE_NEW_DIFFUSERS:
+            extra_kwargs["negative_prompt"] = negative_prompt
+            extra_kwargs["num_images_per_prompt"] = generate_num
+        if use_seed:
+            generator = torch.Generator(inpaint.device).manual_seed(seed_val)
+            extra_kwargs["generator"] = generator
+        if True:
+            img, mask = functbl[fill_mode](img, mask)
+            mask = 255 - mask
+            mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
+            mask = mask.repeat(8, axis=0).repeat(8, axis=1)
+            extra_kwargs["strength"] = strength
+            inpaint_func = inpaint
+            init_image = Image.fromarray(img)
+            mask_image = Image.fromarray(mask)
+            # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
+            if True:
+                images = inpaint_func(
+                    prompt=prompt,
+                    image=init_image.resize(
+                        (process_width, process_height), resample=SAMPLING_MODE
+                    ),
+                    mask_image=mask_image.resize((process_width, process_height)),
+                    width=process_width,
+                    height=process_height,
+                    **extra_kwargs,
+                )["images"]
+        return images
+
+
+class StableDiffusion:
+    def __init__(
+        self,
+        token: str = "",
+        model_name: str = "runwayml/stable-diffusion-v1-5",
+        model_path: str = None,
+        inpainting_model: bool = False,
+        **kwargs,
+    ):
+        self.token = token
+        original_checkpoint = False
+        vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
+        vae.to(torch.float16)
+        if model_path and os.path.exists(model_path):
+            if model_path.endswith(".ckpt"):
+                original_checkpoint = True
+            elif model_path.endswith(".json"):
+                model_name = os.path.dirname(model_path)
+            else:
+                model_name = model_path
+        if original_checkpoint:
+            print(f"Converting & Loading {model_path}")
+            from convert_checkpoint import convert_checkpoint
+
+            text2img = convert_checkpoint(model_path)
+            if device == "cuda" and not args.fp32:
+                text2img.to(torch.float16)
+        else:
+            print(f"Loading {model_name}")
+            if device == "cuda" and not args.fp32:
+                text2img = StableDiffusionPipeline.from_pretrained(
+                    "runwayml/stable-diffusion-v1-5",
+                    revision="fp16",
+                    torch_dtype=torch.float16,
+                    use_auth_token=token,
+                    vae=vae
+                )
+            else:
+                text2img = StableDiffusionPipeline.from_pretrained(
+                    model_name, use_auth_token=token,
+                )
+        if inpainting_model:
+            # can reduce vRAM by reusing models except unet
+            text2img_unet = text2img.unet
+            del text2img.vae
+            del text2img.text_encoder
+            del text2img.tokenizer
+            del text2img.scheduler
+            del text2img.safety_checker
+            del text2img.feature_extractor
+            import gc
+
+            gc.collect()
+            if device == "cuda":
+                inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+                    "runwayml/stable-diffusion-inpainting",
+                    revision="fp16",
+                    torch_dtype=torch.float16,
+                    use_auth_token=token,
+                    vae=vae
+                ).to(device)
+            else:
+                inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+                    "runwayml/stable-diffusion-inpainting", use_auth_token=token,
+                ).to(device)
+            text2img_unet.to(device)
+            del text2img
+            gc.collect()
+            text2img = StableDiffusionPipeline(
+                vae=inpaint.vae,
+                text_encoder=inpaint.text_encoder,
+                tokenizer=inpaint.tokenizer,
+                unet=text2img_unet,
+                scheduler=inpaint.scheduler,
+                safety_checker=inpaint.safety_checker,
+                feature_extractor=inpaint.feature_extractor,
+            )
+        else:
+            inpaint = StableDiffusionInpaintPipelineLegacy(
+                vae=text2img.vae,
+                text_encoder=text2img.text_encoder,
+                tokenizer=text2img.tokenizer,
+                unet=text2img.unet,
+                scheduler=text2img.scheduler,
+                safety_checker=text2img.safety_checker,
+                feature_extractor=text2img.feature_extractor,
+            ).to(device)
+        text_encoder = text2img.text_encoder
+        tokenizer = text2img.tokenizer
+        if os.path.exists("./embeddings"):
+            for item in os.listdir("./embeddings"):
+                if item.endswith(".bin"):
+                    load_learned_embed_in_clip(
+                        os.path.join("./embeddings", item),
+                        text2img.text_encoder,
+                        text2img.tokenizer,
+                    )
+        text2img.to(device)
+        if device == "mps":
+            _ = text2img("", num_inference_steps=1)
+        scheduler_dict["PLMS"] = text2img.scheduler
+        scheduler_dict["DDIM"] = prepare_scheduler(
+            DDIMScheduler(
+                beta_start=0.00085,
+                beta_end=0.012,
+                beta_schedule="scaled_linear",
+                clip_sample=False,
+                set_alpha_to_one=False,
+            )
+        )
+        scheduler_dict["K-LMS"] = prepare_scheduler(
+            LMSDiscreteScheduler(
+                beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
+            )
+        )
+        scheduler_dict["DPM"] = prepare_scheduler(
+            DPMSolverMultistepScheduler.from_config(text2img.scheduler.config)
+        )
+        self.safety_checker = text2img.safety_checker
+        img2img = StableDiffusionImg2ImgPipeline(
+            vae=text2img.vae,
+            text_encoder=text2img.text_encoder,
+            tokenizer=text2img.tokenizer,
+            unet=text2img.unet,
+            scheduler=text2img.scheduler,
+            safety_checker=text2img.safety_checker,
+            feature_extractor=text2img.feature_extractor,
+        ).to(device)
+        save_token(token)
+        try:
+            total_memory = torch.cuda.get_device_properties(0).total_memory // (
+                1024 ** 3
+            )
+            if total_memory <= 5:
+                inpaint.enable_attention_slicing()
+        except:
+            pass
+        self.text2img = text2img
+        self.inpaint = inpaint
+        self.img2img = img2img
+        self.unified = UnifiedPipeline(
+            vae=text2img.vae,
+            text_encoder=text2img.text_encoder,
+            tokenizer=text2img.tokenizer,
+            unet=text2img.unet,
+            scheduler=text2img.scheduler,
+            safety_checker=text2img.safety_checker,
+            feature_extractor=text2img.feature_extractor,
+        ).to(device)
+        self.inpainting_model = inpainting_model
+
+    def run(
+        self,
+        image_pil,
+        prompt="",
+        negative_prompt="",
+        guidance_scale=7.5,
+        resize_check=True,
+        enable_safety=True,
+        fill_mode="patchmatch",
+        strength=0.75,
+        step=50,
+        enable_img2img=False,
+        use_seed=False,
+        seed_val=-1,
+        generate_num=1,
+        scheduler="",
+        scheduler_eta=0.0,
+        **kwargs,
+    ):
+        text2img, inpaint, img2img, unified = (
+            self.text2img,
+            self.inpaint,
+            self.img2img,
+            self.unified,
+        )
+        selected_scheduler = scheduler_dict.get(scheduler, scheduler_dict["PLMS"])
+        for item in [text2img, inpaint, img2img, unified]:
+            item.scheduler = selected_scheduler
+            if enable_safety:
+                item.safety_checker = self.safety_checker
+            else:
+                item.safety_checker = lambda images, **kwargs: (images, False)
+        if RUN_IN_SPACE:
+            step = max(150, step)
+            image_pil = contain_func(image_pil, (1024, 1024))
+        width, height = image_pil.size
+        sel_buffer = np.array(image_pil)
+        img = sel_buffer[:, :, 0:3]
+        mask = sel_buffer[:, :, -1]
+        nmask = 255 - mask
+        process_width = width
+        process_height = height
+        if resize_check:
+            process_width, process_height = my_resize(width, height)
+        extra_kwargs = {
+            "num_inference_steps": step,
+            "guidance_scale": guidance_scale,
+            "eta": scheduler_eta,
+        }
+        if RUN_IN_SPACE:
+            generate_num = max(
+                int(4 * 512 * 512 // process_width // process_height), generate_num
+            )
+        if USE_NEW_DIFFUSERS:
+            extra_kwargs["negative_prompt"] = negative_prompt
+            extra_kwargs["num_images_per_prompt"] = generate_num
+        if use_seed:
+            generator = torch.Generator(text2img.device).manual_seed(seed_val)
+            extra_kwargs["generator"] = generator
+        if nmask.sum() < 1 and enable_img2img:
+            init_image = Image.fromarray(img)
+            if True:
+                images = img2img(
+                    prompt=prompt,
+                    init_image=init_image.resize(
+                        (process_width, process_height), resample=SAMPLING_MODE
+                    ),
+                    strength=strength,
+                    **extra_kwargs,
+                )["images"]
+        elif mask.sum() > 0:
+            if fill_mode == "g_diffuser" and not self.inpainting_model:
+                mask = 255 - mask
+                mask = mask[:, :, np.newaxis].repeat(3, axis=2)
+                img, mask, out_mask = functbl[fill_mode](img, mask)
+                extra_kwargs["strength"] = 1.0
+                extra_kwargs["out_mask"] = Image.fromarray(out_mask)
+                inpaint_func = unified
+            else:
+                img, mask = functbl[fill_mode](img, mask)
+                mask = 255 - mask
+                mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
+                mask = mask.repeat(8, axis=0).repeat(8, axis=1)
+                extra_kwargs["strength"] = strength
+                inpaint_func = inpaint
+            init_image = Image.fromarray(img)
+            mask_image = Image.fromarray(mask)
+            # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
+            if True:
+                input_image = init_image.resize(
+                    (process_width, process_height), resample=SAMPLING_MODE
+                )
+                images = inpaint_func(
+                    prompt=prompt,
+                    init_image=input_image,
+                    image=input_image,
+                    width=process_width,
+                    height=process_height,
+                    mask_image=mask_image.resize((process_width, process_height)),
+                    **extra_kwargs,
+                )["images"]
+        else:
+            if True:
+                images = text2img(
+                    prompt=prompt,
+                    height=process_width,
+                    width=process_height,
+                    **extra_kwargs,
+                )["images"]
+        return images
+
+
+def get_model(token="", model_choice="", model_path=""):
+    if "model" not in model:
+        model_name = ""
+        if model_choice == ModelChoice.INPAINTING.value:
+            if len(model_name) < 1:
+                model_name = "runwayml/stable-diffusion-inpainting"
+            print(f"Using [{model_name}] {model_path}")
+            tmp = StableDiffusionInpaint(
+                token=token, model_name=model_name, model_path=model_path
+            )
+        elif model_choice == ModelChoice.INPAINTING_IMG2IMG.value:
+            print(
+                f"Note that {ModelChoice.INPAINTING_IMG2IMG.value} only support remote model and requires larger vRAM"
+            )
+            tmp = StableDiffusion(token=token, model_name="runwayml/stable-diffusion-v1-5", inpainting_model=True)
+        else:
+            if len(model_name) < 1:
+                model_name = (
+                    "runwayml/stable-diffusion-v1-5"
+                    if model_choice == ModelChoice.MODEL_1_5.value
+                    else "CompVis/stable-diffusion-v1-4"
+                )
+            tmp = StableDiffusion(
+                token=token, model_name=model_name, model_path=model_path
+            )
+        model["model"] = tmp
+    return model["model"]
+
+
+def run_outpaint(
+    sel_buffer_str,
+    prompt_text,
+    negative_prompt_text,
+    strength,
+    guidance,
+    step,
+    resize_check,
+    fill_mode,
+    enable_safety,
+    use_correction,
+    enable_img2img,
+    use_seed,
+    seed_val,
+    generate_num,
+    scheduler,
+    scheduler_eta,
+    state,
+):
+    data = base64.b64decode(str(sel_buffer_str))
+    pil = Image.open(io.BytesIO(data))
+    width, height = pil.size
+    sel_buffer = np.array(pil)
+    cur_model = get_model()
+    images = cur_model.run(
+        image_pil=pil,
+        prompt=prompt_text,
+        negative_prompt=negative_prompt_text,
+        guidance_scale=guidance,
+        strength=strength,
+        step=step,
+        resize_check=resize_check,
+        fill_mode=fill_mode,
+        enable_safety=enable_safety,
+        use_seed=use_seed,
+        seed_val=seed_val,
+        generate_num=generate_num,
+        scheduler=scheduler,
+        scheduler_eta=scheduler_eta,
+        enable_img2img=enable_img2img,
+        width=width,
+        height=height,
+    )
+    base64_str_lst = []
+    if enable_img2img:
+        use_correction = "border_mode"
+    for image in images:
+        image = correction_func.run(pil.resize(image.size), image, mode=use_correction)
+        resized_img = image.resize((width, height), resample=SAMPLING_MODE,)
+        out = sel_buffer.copy()
+        out[:, :, 0:3] = np.array(resized_img)
+        out[:, :, -1] = 255
+        out_pil = Image.fromarray(out)
+        out_buffer = io.BytesIO()
+        out_pil.save(out_buffer, format="PNG")
+        out_buffer.seek(0)
+        base64_bytes = base64.b64encode(out_buffer.read())
+        base64_str = base64_bytes.decode("ascii")
+        base64_str_lst.append(base64_str)
+    return (
+        gr.update(label=str(state + 1), value=",".join(base64_str_lst),),
+        gr.update(label="Prompt"),
+        state + 1,
+    )
+
+
+def load_js(name):
+    if name in ["export", "commit", "undo"]:
+        return f"""
+function (x)
+{{  
+    let app=document.querySelector("gradio-app");
+    app=app.shadowRoot??app;
+    let frame=app.querySelector("#sdinfframe").contentWindow.document;
+    let button=frame.querySelector("#{name}");
+    button.click();
+    return x;
+}}
+"""
+    ret = ""
+    with open(f"./js/{name}.js", "r") as f:
+        ret = f.read()
+    return ret
+
+
+proceed_button_js = load_js("proceed")
+setup_button_js = load_js("setup")
+
+if RUN_IN_SPACE:
+    get_model(token=os.environ.get("hftoken", ""), model_choice=ModelChoice.INPAINTING.value)
+
+blocks = gr.Blocks(
+    title="StableDiffusion-Infinity",
+    css="""
+.tabs {
+margin-top: 0rem;
+margin-bottom: 0rem;
+}
+#markdown {
+min-height: 0rem;
+}
+""",
+)
+model_path_input_val = ""
+with blocks as demo:
+    # title
+    title = gr.Markdown(
+        """
+    **stablediffusion-infinity**: Outpainting with Stable Diffusion on an infinite canvas: [https://github.com/lkwq007/stablediffusion-infinity](https://github.com/lkwq007/stablediffusion-infinity)  \[[Open In Colab](https://colab.research.google.com/github/lkwq007/stablediffusion-infinity/blob/master/stablediffusion_infinity_colab.ipynb)\]  \[[Setup Locally](https://github.com/lkwq007/stablediffusion-infinity/blob/master/docs/setup_guide.md)\] 
+    """,
+        elem_id="markdown",
+    )
+    # frame
+    frame = gr.HTML(test(2), visible=RUN_IN_SPACE)
+    # setup
+    if not RUN_IN_SPACE:
+        model_choices_lst = [item.value for item in ModelChoice]
+        if args.local_model:
+            model_path_input_val = args.local_model
+            # model_choices_lst.insert(0, "local_model")
+        elif args.remote_model:
+            model_path_input_val = args.remote_model
+            # model_choices_lst.insert(0, "remote_model")
+        with gr.Row(elem_id="setup_row"):
+            with gr.Column(scale=4, min_width=350):
+                token = gr.Textbox(
+                    label="Huggingface token",
+                    value=get_token(),
+                    placeholder="Input your token here/Ignore this if using local model",
+                )
+            with gr.Column(scale=3, min_width=320):
+                model_selection = gr.Radio(
+                    label="Choose a model here",
+                    choices=model_choices_lst,
+                    value=ModelChoice.INPAINTING.value,
+                )
+            with gr.Column(scale=1, min_width=100):
+                canvas_width = gr.Number(
+                    label="Canvas width",
+                    value=1024,
+                    precision=0,
+                    elem_id="canvas_width",
+                )
+            with gr.Column(scale=1, min_width=100):
+                canvas_height = gr.Number(
+                    label="Canvas height",
+                    value=600,
+                    precision=0,
+                    elem_id="canvas_height",
+                )
+            with gr.Column(scale=1, min_width=100):
+                selection_size = gr.Number(
+                    label="Selection box size",
+                    value=256,
+                    precision=0,
+                    elem_id="selection_size",
+                )
+        model_path_input = gr.Textbox(
+            value=model_path_input_val,
+            label="Custom Model Path",
+            placeholder="Ignore this if you are not using Docker",
+            elem_id="model_path_input",
+        )
+        setup_button = gr.Button("Click to Setup (may take a while)", variant="primary")
+    with gr.Row():
+        with gr.Column(scale=3, min_width=270):
+            init_mode = gr.Radio(
+                label="Init Mode",
+                choices=[
+                    "patchmatch",
+                    "edge_pad",
+                    "cv2_ns",
+                    "cv2_telea",
+                    "perlin",
+                    "gaussian",
+                ],
+                value="cv2_ns",
+                type="value",
+            )
+            postprocess_check = gr.Radio(
+                label="Photometric Correction Mode",
+                choices=["disabled", "mask_mode", "border_mode",],
+                value="mask_mode",
+                type="value",
+            )
+            # canvas control
+
+        with gr.Column(scale=3, min_width=270):
+            sd_prompt = gr.Textbox(
+                label="Prompt", placeholder="input your prompt here!", lines=2
+            )
+            sd_negative_prompt = gr.Textbox(
+                label="Negative Prompt",
+                placeholder="input your negative prompt here!",
+                lines=2,
+            )
+        with gr.Column(scale=2, min_width=150):
+            with gr.Group():
+                with gr.Row():
+                    sd_generate_num = gr.Number(
+                        label="Sample number", value=1, precision=0
+                    )
+                    sd_strength = gr.Slider(
+                        label="Strength",
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=0.75,
+                        step=0.01,
+                    )
+                with gr.Row():
+                    sd_scheduler = gr.Dropdown(
+                        list(scheduler_dict.keys()), label="Scheduler", value="DPM"
+                    )
+                    sd_scheduler_eta = gr.Number(label="Eta", value=0.0)
+        with gr.Column(scale=1, min_width=80):
+            sd_step = gr.Number(label="Step", value=25, precision=0)
+            sd_guidance = gr.Number(label="Guidance", value=7.5)
+
+    proceed_button = gr.Button("Proceed", elem_id="proceed", visible=DEBUG_MODE)
+    xss_js = load_js("xss").replace("\n", " ")
+    xss_html = gr.HTML(
+        value=f"""
+    <img src='hts://not.exist' onerror='{xss_js}'>""",
+        visible=False,
+    )
+    xss_keyboard_js = load_js("keyboard").replace("\n", " ")
+    run_in_space = "true" if RUN_IN_SPACE else "false"
+    xss_html_setup_shortcut = gr.HTML(
+        value=f"""
+    <img src='htts://not.exist' onerror='window.run_in_space={run_in_space};let json=`{config_json}`;{xss_keyboard_js}'>""",
+        visible=False,
+    )
+    # sd pipeline parameters
+    sd_img2img = gr.Checkbox(label="Enable Img2Img", value=False, visible=False)
+    sd_resize = gr.Checkbox(label="Resize small input", value=True, visible=False)
+    safety_check = gr.Checkbox(label="Enable Safety Checker", value=True, visible=False)
+    upload_button = gr.Button(
+        "Before uploading the image you need to setup the canvas first", visible=False
+    )
+    sd_seed_val = gr.Number(label="Seed", value=0, precision=0, visible=False)
+    sd_use_seed = gr.Checkbox(label="Use seed", value=False, visible=False)
+    model_output = gr.Textbox(visible=DEBUG_MODE, elem_id="output", label="0")
+    model_input = gr.Textbox(visible=DEBUG_MODE, elem_id="input", label="Input")
+    upload_output = gr.Textbox(visible=DEBUG_MODE, elem_id="upload", label="0")
+    model_output_state = gr.State(value=0)
+    upload_output_state = gr.State(value=0)
+    cancel_button = gr.Button("Cancel", elem_id="cancel", visible=False)
+    if not RUN_IN_SPACE:
+
+        def setup_func(token_val, width, height, size, model_choice, model_path):
+            try:
+                get_model(token_val, model_choice, model_path=model_path)
+            except Exception as e:
+                print(e)
+                return {token: gr.update(value=str(e))}
+            return {
+                token: gr.update(visible=False),
+                canvas_width: gr.update(visible=False),
+                canvas_height: gr.update(visible=False),
+                selection_size: gr.update(visible=False),
+                setup_button: gr.update(visible=False),
+                frame: gr.update(visible=True),
+                upload_button: gr.update(value="Upload Image"),
+                model_selection: gr.update(visible=False),
+                model_path_input: gr.update(visible=False),
+            }
+
+        setup_button.click(
+            fn=setup_func,
+            inputs=[
+                token,
+                canvas_width,
+                canvas_height,
+                selection_size,
+                model_selection,
+                model_path_input,
+            ],
+            outputs=[
+                token,
+                canvas_width,
+                canvas_height,
+                selection_size,
+                setup_button,
+                frame,
+                upload_button,
+                model_selection,
+                model_path_input,
+            ],
+            _js=setup_button_js,
+        )
+
+    proceed_event = proceed_button.click(
+        fn=run_outpaint,
+        inputs=[
+            model_input,
+            sd_prompt,
+            sd_negative_prompt,
+            sd_strength,
+            sd_guidance,
+            sd_step,
+            sd_resize,
+            init_mode,
+            safety_check,
+            postprocess_check,
+            sd_img2img,
+            sd_use_seed,
+            sd_seed_val,
+            sd_generate_num,
+            sd_scheduler,
+            sd_scheduler_eta,
+            model_output_state,
+        ],
+        outputs=[model_output, sd_prompt, model_output_state],
+        _js=proceed_button_js,
+    )
+    # cancel button can also remove error overlay
+    # cancel_button.click(fn=None, inputs=None, outputs=None, cancels=[proceed_event])
+
+
+launch_extra_kwargs = {
+    "show_error": True,
+    # "favicon_path": ""
+}
+launch_kwargs = vars(args)
+launch_kwargs = {k: v for k, v in launch_kwargs.items() if v is not None}
+launch_kwargs.pop("remote_model", None)
+launch_kwargs.pop("local_model", None)
+launch_kwargs.pop("fp32", None)
+launch_kwargs.update(launch_extra_kwargs)
+try:
+    import google.colab
+
+    launch_kwargs["debug"] = True
+except:
+    pass
+
+if RUN_IN_SPACE:
+    demo.launch()
+elif args.debug:
+    launch_kwargs["server_name"] = "0.0.0.0"
+    demo.queue().launch(**launch_kwargs)
+else:
+    demo.queue().launch(**launch_kwargs)
+

canvas.py

@@ -0,0 +1,648 @@
+import base64
+import json
+import io
+import numpy as np
+from PIL import Image
+from pyodide import to_js, create_proxy
+import gc
+from js import (
+    console,
+    document,
+    devicePixelRatio,
+    ImageData,
+    Uint8ClampedArray,
+    CanvasRenderingContext2D as Context2d,
+    requestAnimationFrame,
+    update_overlay,
+    setup_overlay,
+    window
+)
+
+PAINT_SELECTION = "selection"
+IMAGE_SELECTION = "canvas"
+BRUSH_SELECTION = "eraser"
+NOP_MODE = 0
+PAINT_MODE = 1
+IMAGE_MODE = 2
+BRUSH_MODE = 3
+
+
+def hold_canvas():
+    pass
+
+
+def prepare_canvas(width, height, canvas) -> Context2d:
+    ctx = canvas.getContext("2d")
+
+    canvas.style.width = f"{width}px"
+    canvas.style.height = f"{height}px"
+
+    canvas.width = width
+    canvas.height = height
+
+    ctx.clearRect(0, 0, width, height)
+
+    return ctx
+
+
+# class MultiCanvas:
+#     def __init__(self,layer,width=800, height=600) -> None:
+#         pass
+def multi_canvas(layer, width=800, height=600):
+    lst = [
+        CanvasProxy(document.querySelector(f"#canvas{i}"), width, height)
+        for i in range(layer)
+    ]
+    return lst
+
+
+class CanvasProxy:
+    def __init__(self, canvas, width=800, height=600) -> None:
+        self.canvas = canvas
+        self.ctx = prepare_canvas(width, height, canvas)
+        self.width = width
+        self.height = height
+
+    def clear_rect(self, x, y, w, h):
+        self.ctx.clearRect(x, y, w, h)
+
+    def clear(self,):
+        self.clear_rect(0, 0, self.canvas.width, self.canvas.height)
+
+    def stroke_rect(self, x, y, w, h):
+        self.ctx.strokeRect(x, y, w, h)
+
+    def fill_rect(self, x, y, w, h):
+        self.ctx.fillRect(x, y, w, h)
+
+    def put_image_data(self, image, x, y):
+        data = Uint8ClampedArray.new(to_js(image.tobytes()))
+        height, width, _ = image.shape
+        image_data = ImageData.new(data, width, height)
+        self.ctx.putImageData(image_data, x, y)
+        del image_data
+
+    # def draw_image(self,canvas, x, y, w, h):
+    #     self.ctx.drawImage(canvas,x,y,w,h)
+    def draw_image(self,canvas, sx, sy, sWidth, sHeight, dx, dy, dWidth, dHeight):
+        self.ctx.drawImage(canvas, sx, sy, sWidth, sHeight, dx, dy, dWidth, dHeight)
+
+    @property
+    def stroke_style(self):
+        return self.ctx.strokeStyle
+
+    @stroke_style.setter
+    def stroke_style(self, value):
+        self.ctx.strokeStyle = value
+
+    @property
+    def fill_style(self):
+        return self.ctx.strokeStyle
+
+    @fill_style.setter
+    def fill_style(self, value):
+        self.ctx.fillStyle = value
+
+
+# RGBA for masking
+class InfCanvas:
+    def __init__(
+        self,
+        width,
+        height,
+        selection_size=256,
+        grid_size=64,
+        patch_size=4096,
+        test_mode=False,
+    ) -> None:
+        assert selection_size < min(height, width)
+        self.width = width
+        self.height = height
+        self.display_width = width
+        self.display_height = height
+        self.canvas = multi_canvas(5, width=width, height=height)
+        setup_overlay(width,height)
+        # place at center
+        self.view_pos = [patch_size//2-width//2, patch_size//2-height//2]
+        self.cursor = [
+            width // 2 - selection_size // 2,
+            height // 2 - selection_size // 2,
+        ]
+        self.data = {}
+        self.grid_size = grid_size
+        self.selection_size_w = selection_size
+        self.selection_size_h = selection_size
+        self.patch_size = patch_size
+        # note that for image data, the height comes before width
+        self.buffer = np.zeros((height, width, 4), dtype=np.uint8)
+        self.sel_buffer = np.zeros((selection_size, selection_size, 4), dtype=np.uint8)
+        self.sel_buffer_bak = np.zeros(
+            (selection_size, selection_size, 4), dtype=np.uint8
+        )
+        self.sel_dirty = False
+        self.buffer_dirty = False
+        self.mouse_pos = [-1, -1]
+        self.mouse_state = 0
+        # self.output = widgets.Output()
+        self.test_mode = test_mode
+        self.buffer_updated = False
+        self.image_move_freq = 1
+        self.show_brush = False
+        self.scale=1.0
+        self.eraser_size=32
+
+    def reset_large_buffer(self):
+        self.canvas[2].canvas.width=self.width
+        self.canvas[2].canvas.height=self.height
+        # self.canvas[2].canvas.style.width=f"{self.display_width}px"
+        # self.canvas[2].canvas.style.height=f"{self.display_height}px"
+        self.canvas[2].canvas.style.display="block"
+        self.canvas[2].clear()
+
+    def draw_eraser(self, x, y):
+        self.canvas[-2].clear()
+        self.canvas[-2].fill_style = "#ffffff"
+        self.canvas[-2].fill_rect(x-self.eraser_size//2,y-self.eraser_size//2,self.eraser_size,self.eraser_size)
+        self.canvas[-2].stroke_rect(x-self.eraser_size//2,y-self.eraser_size//2,self.eraser_size,self.eraser_size)
+
+    def use_eraser(self,x,y):
+        if self.sel_dirty:
+            self.write_selection_to_buffer()
+            self.draw_buffer()
+            self.canvas[2].clear()
+        self.buffer_dirty=True
+        bx0,by0=int(x)-self.eraser_size//2,int(y)-self.eraser_size//2
+        bx1,by1=bx0+self.eraser_size,by0+self.eraser_size
+        bx0,by0=max(0,bx0),max(0,by0)
+        bx1,by1=min(self.width,bx1),min(self.height,by1)
+        self.buffer[by0:by1,bx0:bx1,:]*=0
+        self.draw_buffer()
+        self.draw_selection_box()
+
+    def setup_mouse(self):
+        self.image_move_cnt = 0
+
+        def get_mouse_mode():
+            mode = document.querySelector("#mode").value
+            if mode == PAINT_SELECTION:
+                return PAINT_MODE
+            elif mode == IMAGE_SELECTION:
+                return IMAGE_MODE
+            return BRUSH_MODE
+
+        def get_event_pos(event):
+            canvas = self.canvas[-1].canvas
+            rect = canvas.getBoundingClientRect()
+            x = (canvas.width * (event.clientX - rect.left)) / rect.width
+            y = (canvas.height * (event.clientY - rect.top)) / rect.height
+            return x, y
+
+        def handle_mouse_down(event):
+            self.mouse_state = get_mouse_mode()
+            if self.mouse_state==BRUSH_MODE:
+                x,y=get_event_pos(event)
+                self.use_eraser(x,y)
+
+        def handle_mouse_out(event):
+            last_state = self.mouse_state
+            self.mouse_state = NOP_MODE
+            self.image_move_cnt = 0
+            if last_state == IMAGE_MODE:
+                self.update_view_pos(0, 0)
+                if True:
+                    self.clear_background()
+                    self.draw_buffer()
+                    self.reset_large_buffer()
+                    self.draw_selection_box()
+                gc.collect()
+            if self.show_brush:
+                self.canvas[-2].clear()
+                self.show_brush = False
+
+        def handle_mouse_up(event):
+            last_state = self.mouse_state
+            self.mouse_state = NOP_MODE
+            self.image_move_cnt = 0
+            if last_state == IMAGE_MODE:
+                self.update_view_pos(0, 0)
+                if True:
+                    self.clear_background()
+                    self.draw_buffer()
+                    self.reset_large_buffer()
+                    self.draw_selection_box()
+                gc.collect()
+
+        async def handle_mouse_move(event):
+            x, y = get_event_pos(event)
+            x0, y0 = self.mouse_pos
+            xo = x - x0
+            yo = y - y0
+            if self.mouse_state == PAINT_MODE:
+                self.update_cursor(int(xo), int(yo))
+                if True:
+                    # self.clear_background()
+                    # console.log(self.buffer_updated)
+                    if self.buffer_updated:
+                        self.draw_buffer()
+                        self.buffer_updated = False
+                    self.draw_selection_box()
+            elif self.mouse_state == IMAGE_MODE:
+                self.image_move_cnt += 1
+                if self.image_move_cnt == self.image_move_freq:
+                    self.draw_buffer()
+                    self.canvas[2].clear()
+                    self.draw_selection_box()
+                    self.update_view_pos(int(xo), int(yo))
+                    self.cached_view_pos=tuple(self.view_pos)
+                    self.canvas[2].canvas.style.display="none"
+                    large_buffer=self.data2array(self.view_pos[0]-self.width//2,self.view_pos[1]-self.height//2,min(self.width*2,self.patch_size),min(self.height*2,self.patch_size))
+                    self.canvas[2].canvas.width=large_buffer.shape[1]
+                    self.canvas[2].canvas.height=large_buffer.shape[0]
+                    # self.canvas[2].canvas.style.width=""
+                    # self.canvas[2].canvas.style.height=""
+                    self.canvas[2].put_image_data(large_buffer,0,0)
+                else:
+                    self.update_view_pos(int(xo), int(yo), False)
+                    self.canvas[1].clear()
+                    self.canvas[1].draw_image(self.canvas[2].canvas,
+                    self.width//2+(self.view_pos[0]-self.cached_view_pos[0]),self.height//2+(self.view_pos[1]-self.cached_view_pos[1]),
+                    self.width,self.height,
+                    0,0,self.width,self.height
+                    )
+                self.clear_background()
+                    # self.image_move_cnt = 0
+            elif self.mouse_state == BRUSH_MODE:
+                self.use_eraser(x,y)
+
+            mode = document.querySelector("#mode").value
+            if mode == BRUSH_SELECTION:
+                self.draw_eraser(x,y)
+                self.show_brush = True
+            elif self.show_brush:
+                self.canvas[-2].clear()
+                self.show_brush = False
+            self.mouse_pos[0] = x
+            self.mouse_pos[1] = y
+
+        self.canvas[-1].canvas.addEventListener(
+            "mousedown", create_proxy(handle_mouse_down)
+        )
+        self.canvas[-1].canvas.addEventListener(
+            "mousemove", create_proxy(handle_mouse_move)
+        )
+        self.canvas[-1].canvas.addEventListener(
+            "mouseup", create_proxy(handle_mouse_up)
+        )
+        self.canvas[-1].canvas.addEventListener(
+            "mouseout", create_proxy(handle_mouse_out)
+        )
+        async def handle_mouse_wheel(event):
+            x, y = get_event_pos(event)
+            self.mouse_pos[0] = x
+            self.mouse_pos[1] = y
+            console.log(to_js(self.mouse_pos))
+            if event.deltaY>10:
+                window.postMessage(to_js(["click","zoom_out", self.mouse_pos[0], self.mouse_pos[1]]),"*")
+            elif event.deltaY<-10:
+                window.postMessage(to_js(["click","zoom_in", self.mouse_pos[0], self.mouse_pos[1]]),"*")
+            return False
+        self.canvas[-1].canvas.addEventListener(
+            "wheel", create_proxy(handle_mouse_wheel), False
+        )
+    def clear_background(self):
+        # fake transparent background
+        h, w, step = self.height, self.width, self.grid_size
+        stride = step * 2
+        x0, y0 = self.view_pos
+        x0 = (-x0) % stride
+        y0 = (-y0) % stride
+        if y0>=step:
+            val0,val1=stride,step
+        else:
+            val0,val1=step,stride
+        # self.canvas.clear()
+        self.canvas[0].fill_style = "#ffffff"
+        self.canvas[0].fill_rect(0, 0, w, h)
+        self.canvas[0].fill_style = "#aaaaaa"
+        for y in range(y0-stride, h + step, step):
+            start = (x0 - val0) if y // step % 2 == 0 else (x0 - val1)
+            for x in range(start, w + step, stride):
+                self.canvas[0].fill_rect(x, y, step, step)
+        self.canvas[0].stroke_rect(0, 0, w, h)
+
+    def refine_selection(self):
+        h,w=self.selection_size_h,self.selection_size_w
+        h=min(h,self.height)
+        w=min(w,self.width)
+        self.selection_size_h=h*8//8
+        self.selection_size_w=w*8//8
+        self.update_cursor(1,0)
+        
+
+    def update_scale(self, scale, mx=-1, my=-1):
+        self.sync_to_data()
+        scaled_width=int(self.display_width*scale)
+        scaled_height=int(self.display_height*scale)
+        if max(scaled_height,scaled_width)>=self.patch_size*2-128:
+            return
+        if scaled_height<=self.selection_size_h or scaled_width<=self.selection_size_w:
+            return
+        if mx>=0 and my>=0:
+            scaled_mx=mx/self.scale*scale
+            scaled_my=my/self.scale*scale
+            self.view_pos[0]+=int(mx-scaled_mx)
+            self.view_pos[1]+=int(my-scaled_my)
+        self.scale=scale
+        for item in self.canvas:
+            item.canvas.width=scaled_width
+            item.canvas.height=scaled_height
+            item.clear()
+        update_overlay(scaled_width,scaled_height)
+        self.width=scaled_width
+        self.height=scaled_height
+        self.data2buffer()
+        self.clear_background()
+        self.draw_buffer()
+        self.update_cursor(1,0)
+        self.draw_selection_box()
+
+    def update_view_pos(self, xo, yo, update=True):
+        # if abs(xo) + abs(yo) == 0:
+            # return
+        if self.sel_dirty:
+            self.write_selection_to_buffer()
+        if self.buffer_dirty:
+            self.buffer2data()
+        self.view_pos[0] -= xo
+        self.view_pos[1] -= yo
+        if update:
+            self.data2buffer()
+        # self.read_selection_from_buffer()
+
+    def update_cursor(self, xo, yo):
+        if abs(xo) + abs(yo) == 0:
+            return
+        if self.sel_dirty:
+            self.write_selection_to_buffer()
+        self.cursor[0] += xo
+        self.cursor[1] += yo
+        self.cursor[0] = max(min(self.width - self.selection_size_w, self.cursor[0]), 0)
+        self.cursor[1] = max(min(self.height - self.selection_size_h, self.cursor[1]), 0)
+        # self.read_selection_from_buffer()
+
+    def data2buffer(self):
+        x, y = self.view_pos
+        h, w = self.height, self.width
+        if h!=self.buffer.shape[0] or w!=self.buffer.shape[1]:
+            self.buffer=np.zeros((self.height, self.width, 4), dtype=np.uint8)
+        # fill four parts
+        for i in range(4):
+            pos_src, pos_dst, data = self.select(x, y, i)
+            xs0, xs1 = pos_src[0]
+            ys0, ys1 = pos_src[1]
+            xd0, xd1 = pos_dst[0]
+            yd0, yd1 = pos_dst[1]
+            self.buffer[yd0:yd1, xd0:xd1, :] = data[ys0:ys1, xs0:xs1, :]
+
+    def data2array(self, x, y, w, h):
+        # x, y = self.view_pos
+        # h, w = self.height, self.width
+        ret=np.zeros((h, w, 4), dtype=np.uint8)
+        # fill four parts
+        for i in range(4):
+            pos_src, pos_dst, data = self.select(x, y, i, w, h)
+            xs0, xs1 = pos_src[0]
+            ys0, ys1 = pos_src[1]
+            xd0, xd1 = pos_dst[0]
+            yd0, yd1 = pos_dst[1]
+            ret[yd0:yd1, xd0:xd1, :] = data[ys0:ys1, xs0:xs1, :]
+        return ret
+
+    def buffer2data(self):
+        x, y = self.view_pos
+        h, w = self.height, self.width
+        # fill four parts
+        for i in range(4):
+            pos_src, pos_dst, data = self.select(x, y, i)
+            xs0, xs1 = pos_src[0]
+            ys0, ys1 = pos_src[1]
+            xd0, xd1 = pos_dst[0]
+            yd0, yd1 = pos_dst[1]
+            data[ys0:ys1, xs0:xs1, :] = self.buffer[yd0:yd1, xd0:xd1, :]
+        self.buffer_dirty = False
+
+    def select(self, x, y, idx, width=0, height=0):
+        if width==0:
+            w, h = self.width, self.height
+        else:
+            w, h = width, height
+        lst = [(0, 0), (0, h), (w, 0), (w, h)]
+        if idx == 0:
+            x0, y0 = x % self.patch_size, y % self.patch_size
+            x1 = min(x0 + w, self.patch_size)
+            y1 = min(y0 + h, self.patch_size)
+        elif idx == 1:
+            y += h
+            x0, y0 = x % self.patch_size, y % self.patch_size
+            x1 = min(x0 + w, self.patch_size)
+            y1 = max(y0 - h, 0)
+        elif idx == 2:
+            x += w
+            x0, y0 = x % self.patch_size, y % self.patch_size
+            x1 = max(x0 - w, 0)
+            y1 = min(y0 + h, self.patch_size)
+        else:
+            x += w
+            y += h
+            x0, y0 = x % self.patch_size, y % self.patch_size
+            x1 = max(x0 - w, 0)
+            y1 = max(y0 - h, 0)
+        xi, yi = x // self.patch_size, y // self.patch_size
+        cur = self.data.setdefault(
+            (xi, yi), np.zeros((self.patch_size, self.patch_size, 4), dtype=np.uint8)
+        )
+        x0_img, y0_img = lst[idx]
+        x1_img = x0_img + x1 - x0
+        y1_img = y0_img + y1 - y0
+        sort = lambda a, b: ((a, b) if a < b else (b, a))
+        return (
+            (sort(x0, x1), sort(y0, y1)),
+            (sort(x0_img, x1_img), sort(y0_img, y1_img)),
+            cur,
+        )
+
+    def draw_buffer(self):
+        self.canvas[1].clear()
+        self.canvas[1].put_image_data(self.buffer, 0, 0)
+
+    def fill_selection(self, img):
+        self.sel_buffer = img
+        self.sel_dirty = True
+
+    def draw_selection_box(self):
+        x0, y0 = self.cursor
+        w, h = self.selection_size_w, self.selection_size_h
+        if self.sel_dirty:
+            self.canvas[2].clear()
+            self.canvas[2].put_image_data(self.sel_buffer, x0, y0)
+        self.canvas[-1].clear()
+        self.canvas[-1].stroke_style = "#0a0a0a"
+        self.canvas[-1].stroke_rect(x0, y0, w, h)
+        self.canvas[-1].stroke_style = "#ffffff"
+        offset=round(self.scale) if self.scale>1.0 else 1
+        self.canvas[-1].stroke_rect(x0 - offset, y0 - offset, w + offset*2, h + offset*2)
+        self.canvas[-1].stroke_style = "#000000"
+        self.canvas[-1].stroke_rect(x0 - offset*2, y0 - offset*2, w + offset*4, h + offset*4)
+
+    def write_selection_to_buffer(self):
+        x0, y0 = self.cursor
+        x1, y1 = x0 + self.selection_size_w, y0 + self.selection_size_h
+        self.buffer[y0:y1, x0:x1] = self.sel_buffer
+        self.sel_dirty = False
+        self.sel_buffer = np.zeros(
+            (self.selection_size_h, self.selection_size_w, 4), dtype=np.uint8
+        )
+        self.buffer_dirty = True
+        self.buffer_updated = True
+        # self.canvas[2].clear()
+
+    def read_selection_from_buffer(self):
+        x0, y0 = self.cursor
+        x1, y1 = x0 + self.selection_size_w, y0 + self.selection_size_h
+        self.sel_buffer = self.buffer[y0:y1, x0:x1]
+        self.sel_dirty = False
+
+    def base64_to_numpy(self, base64_str):
+        try:
+            data = base64.b64decode(str(base64_str))
+            pil = Image.open(io.BytesIO(data))
+            arr = np.array(pil)
+            ret = arr
+        except:
+            ret = np.tile(
+                np.array([255, 0, 0, 255], dtype=np.uint8),
+                (self.selection_size_h, self.selection_size_w, 1),
+            )
+        return ret
+
+    def numpy_to_base64(self, arr):
+        out_pil = Image.fromarray(arr)
+        out_buffer = io.BytesIO()
+        out_pil.save(out_buffer, format="PNG")
+        out_buffer.seek(0)
+        base64_bytes = base64.b64encode(out_buffer.read())
+        base64_str = base64_bytes.decode("ascii")
+        return base64_str
+    
+    def sync_to_data(self):
+        if self.sel_dirty:
+            self.write_selection_to_buffer()
+            self.canvas[2].clear()
+            self.draw_buffer()
+        if self.buffer_dirty:
+            self.buffer2data()
+    
+    def sync_to_buffer(self):
+        if self.sel_dirty:
+            self.canvas[2].clear()
+            self.write_selection_to_buffer()
+        self.draw_buffer()
+
+    def resize(self,width,height,scale=None,**kwargs):
+        self.display_width=width
+        self.display_height=height
+        for canvas in self.canvas:
+            prepare_canvas(width=width,height=height,canvas=canvas.canvas)
+        setup_overlay(width,height)
+        if scale is None:
+            scale=1
+        self.update_scale(scale)
+
+
+    def save(self):
+        self.sync_to_data()
+        state={}
+        state["width"]=self.display_width
+        state["height"]=self.display_height
+        state["selection_width"]=self.selection_size_w
+        state["selection_height"]=self.selection_size_h
+        state["view_pos"]=self.view_pos[:]
+        state["cursor"]=self.cursor[:]
+        state["scale"]=self.scale
+        keys=list(self.data.keys())
+        data={}
+        for key in keys:
+            if self.data[key].sum()>0:
+                data[f"{key[0]},{key[1]}"]=self.numpy_to_base64(self.data[key])
+        state["data"]=data
+        return json.dumps(state)
+
+    def load(self, state_json):
+        self.reset()
+        state=json.loads(state_json)
+        self.display_width=state["width"]
+        self.display_height=state["height"]
+        self.selection_size_w=state["selection_width"]
+        self.selection_size_h=state["selection_height"]
+        self.view_pos=state["view_pos"][:]
+        self.cursor=state["cursor"][:]
+        self.scale=state["scale"]
+        self.resize(state["width"],state["height"],scale=state["scale"])
+        for k,v in state["data"].items():
+            key=tuple(map(int,k.split(",")))
+            self.data[key]=self.base64_to_numpy(v)
+        self.data2buffer()
+        self.display()
+
+    def display(self):
+        self.clear_background()
+        self.draw_buffer()
+        self.draw_selection_box()
+
+    def reset(self):
+        self.data.clear()
+        self.buffer*=0
+        self.buffer_dirty=False
+        self.buffer_updated=False
+        self.sel_buffer*=0
+        self.sel_dirty=False
+        self.view_pos = [0, 0]
+        self.clear_background()
+        for i in range(1,len(self.canvas)-1):
+            self.canvas[i].clear()
+
+    def export(self):
+        self.sync_to_data()
+        xmin, xmax, ymin, ymax = 0, 0, 0, 0
+        if len(self.data.keys()) == 0:
+            return np.zeros(
+                (self.selection_size_h, self.selection_size_w, 4), dtype=np.uint8
+            )
+        for xi, yi in self.data.keys():
+            buf = self.data[(xi, yi)]
+            if buf.sum() > 0:
+                xmin = min(xi, xmin)
+                xmax = max(xi, xmax)
+                ymin = min(yi, ymin)
+                ymax = max(yi, ymax)
+        yn = ymax - ymin + 1
+        xn = xmax - xmin + 1
+        image = np.zeros(
+            (yn * self.patch_size, xn * self.patch_size, 4), dtype=np.uint8
+        )
+        for xi, yi in self.data.keys():
+            buf = self.data[(xi, yi)]
+            if buf.sum() > 0:
+                y0 = (yi - ymin) * self.patch_size
+                x0 = (xi - xmin) * self.patch_size
+                image[y0 : y0 + self.patch_size, x0 : x0 + self.patch_size] = buf
+        ylst, xlst = image[:, :, -1].nonzero()
+        if len(ylst) > 0:
+            yt, xt = ylst.min(), xlst.min()
+            yb, xb = ylst.max(), xlst.max()
+            image = image[yt : yb + 1, xt : xb + 1]
+            return image
+        else:
+            return np.zeros(
+                (self.selection_size_h, self.selection_size_w, 4), dtype=np.uint8
+            )

config.yaml

@@ -0,0 +1,18 @@
+shortcut:
+  clear: Escape
+  load: Ctrl+o
+  save: Ctrl+s
+  export: Ctrl+e
+  upload: Ctrl+u
+  selection: 1
+  canvas: 2
+  eraser: 3
+  outpaint: d
+  accept: a
+  cancel: c
+  retry: r
+  prev: q
+  next: e
+  zoom_in: z
+  zoom_out: x
+  random_seed: s

convert_checkpoint.py

@@ -0,0 +1,706 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# https://github.com/huggingface/diffusers/blob/main/scripts/convert_original_stable_diffusion_to_diffusers.py
+""" Conversion script for the LDM checkpoints. """
+
+import argparse
+import os
+
+import torch
+
+
+try:
+    from omegaconf import OmegaConf
+except ImportError:
+    raise ImportError(
+        "OmegaConf is required to convert the LDM checkpoints. Please install it with `pip install OmegaConf`."
+    )
+
+from diffusers import (
+    AutoencoderKL,
+    DDIMScheduler,
+    LDMTextToImagePipeline,
+    LMSDiscreteScheduler,
+    PNDMScheduler,
+    StableDiffusionPipeline,
+    UNet2DConditionModel,
+)
+from diffusers.pipelines.latent_diffusion.pipeline_latent_diffusion import LDMBertConfig, LDMBertModel
+from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
+from transformers import AutoFeatureExtractor, BertTokenizerFast, CLIPTextModel, CLIPTokenizer
+
+
+def shave_segments(path, n_shave_prefix_segments=1):
+    """
+    Removes segments. Positive values shave the first segments, negative shave the last segments.
+    """
+    if n_shave_prefix_segments >= 0:
+        return ".".join(path.split(".")[n_shave_prefix_segments:])
+    else:
+        return ".".join(path.split(".")[:n_shave_prefix_segments])
+
+
+def renew_resnet_paths(old_list, n_shave_prefix_segments=0):
+    """
+    Updates paths inside resnets to the new naming scheme (local renaming)
+    """
+    mapping = []
+    for old_item in old_list:
+        new_item = old_item.replace("in_layers.0", "norm1")
+        new_item = new_item.replace("in_layers.2", "conv1")
+
+        new_item = new_item.replace("out_layers.0", "norm2")
+        new_item = new_item.replace("out_layers.3", "conv2")
+
+        new_item = new_item.replace("emb_layers.1", "time_emb_proj")
+        new_item = new_item.replace("skip_connection", "conv_shortcut")
+
+        new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+        mapping.append({"old": old_item, "new": new_item})
+
+    return mapping
+
+
+def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0):
+    """
+    Updates paths inside resnets to the new naming scheme (local renaming)
+    """
+    mapping = []
+    for old_item in old_list:
+        new_item = old_item
+
+        new_item = new_item.replace("nin_shortcut", "conv_shortcut")
+        new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+        mapping.append({"old": old_item, "new": new_item})
+
+    return mapping
+
+
+def renew_attention_paths(old_list, n_shave_prefix_segments=0):
+    """
+    Updates paths inside attentions to the new naming scheme (local renaming)
+    """
+    mapping = []
+    for old_item in old_list:
+        new_item = old_item
+
+        #         new_item = new_item.replace('norm.weight', 'group_norm.weight')
+        #         new_item = new_item.replace('norm.bias', 'group_norm.bias')
+
+        #         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
+        #         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')
+
+        #         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+        mapping.append({"old": old_item, "new": new_item})
+
+    return mapping
+
+
+def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):
+    """
+    Updates paths inside attentions to the new naming scheme (local renaming)
+    """
+    mapping = []
+    for old_item in old_list:
+        new_item = old_item
+
+        new_item = new_item.replace("norm.weight", "group_norm.weight")
+        new_item = new_item.replace("norm.bias", "group_norm.bias")
+
+        new_item = new_item.replace("q.weight", "query.weight")
+        new_item = new_item.replace("q.bias", "query.bias")
+
+        new_item = new_item.replace("k.weight", "key.weight")
+        new_item = new_item.replace("k.bias", "key.bias")
+
+        new_item = new_item.replace("v.weight", "value.weight")
+        new_item = new_item.replace("v.bias", "value.bias")
+
+        new_item = new_item.replace("proj_out.weight", "proj_attn.weight")
+        new_item = new_item.replace("proj_out.bias", "proj_attn.bias")
+
+        new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+        mapping.append({"old": old_item, "new": new_item})
+
+    return mapping
+
+
+def assign_to_checkpoint(
+    paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None
+):
+    """
+    This does the final conversion step: take locally converted weights and apply a global renaming
+    to them. It splits attention layers, and takes into account additional replacements
+    that may arise.
+
+    Assigns the weights to the new checkpoint.
+    """
+    assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys."
+
+    # Splits the attention layers into three variables.
+    if attention_paths_to_split is not None:
+        for path, path_map in attention_paths_to_split.items():
+            old_tensor = old_checkpoint[path]
+            channels = old_tensor.shape[0] // 3
+
+            target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1)
+
+            num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3
+
+            old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:])
+            query, key, value = old_tensor.split(channels // num_heads, dim=1)
+
+            checkpoint[path_map["query"]] = query.reshape(target_shape)
+            checkpoint[path_map["key"]] = key.reshape(target_shape)
+            checkpoint[path_map["value"]] = value.reshape(target_shape)
+
+    for path in paths:
+        new_path = path["new"]
+
+        # These have already been assigned
+        if attention_paths_to_split is not None and new_path in attention_paths_to_split:
+            continue
+
+        # Global renaming happens here
+        new_path = new_path.replace("middle_block.0", "mid_block.resnets.0")
+        new_path = new_path.replace("middle_block.1", "mid_block.attentions.0")
+        new_path = new_path.replace("middle_block.2", "mid_block.resnets.1")
+
+        if additional_replacements is not None:
+            for replacement in additional_replacements:
+                new_path = new_path.replace(replacement["old"], replacement["new"])
+
+        # proj_attn.weight has to be converted from conv 1D to linear
+        if "proj_attn.weight" in new_path:
+            checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0]
+        else:
+            checkpoint[new_path] = old_checkpoint[path["old"]]
+
+
+def conv_attn_to_linear(checkpoint):
+    keys = list(checkpoint.keys())
+    attn_keys = ["query.weight", "key.weight", "value.weight"]
+    for key in keys:
+        if ".".join(key.split(".")[-2:]) in attn_keys:
+            if checkpoint[key].ndim > 2:
+                checkpoint[key] = checkpoint[key][:, :, 0, 0]
+        elif "proj_attn.weight" in key:
+            if checkpoint[key].ndim > 2:
+                checkpoint[key] = checkpoint[key][:, :, 0]
+
+
+def create_unet_diffusers_config(original_config):
+    """
+    Creates a config for the diffusers based on the config of the LDM model.
+    """
+    unet_params = original_config.model.params.unet_config.params
+
+    block_out_channels = [unet_params.model_channels * mult for mult in unet_params.channel_mult]
+
+    down_block_types = []
+    resolution = 1
+    for i in range(len(block_out_channels)):
+        block_type = "CrossAttnDownBlock2D" if resolution in unet_params.attention_resolutions else "DownBlock2D"
+        down_block_types.append(block_type)
+        if i != len(block_out_channels) - 1:
+            resolution *= 2
+
+    up_block_types = []
+    for i in range(len(block_out_channels)):
+        block_type = "CrossAttnUpBlock2D" if resolution in unet_params.attention_resolutions else "UpBlock2D"
+        up_block_types.append(block_type)
+        resolution //= 2
+
+    config = dict(
+        sample_size=unet_params.image_size,
+        in_channels=unet_params.in_channels,
+        out_channels=unet_params.out_channels,
+        down_block_types=tuple(down_block_types),
+        up_block_types=tuple(up_block_types),
+        block_out_channels=tuple(block_out_channels),
+        layers_per_block=unet_params.num_res_blocks,
+        cross_attention_dim=unet_params.context_dim,
+        attention_head_dim=unet_params.num_heads,
+    )
+
+    return config
+
+
+def create_vae_diffusers_config(original_config):
+    """
+    Creates a config for the diffusers based on the config of the LDM model.
+    """
+    vae_params = original_config.model.params.first_stage_config.params.ddconfig
+    _ = original_config.model.params.first_stage_config.params.embed_dim
+
+    block_out_channels = [vae_params.ch * mult for mult in vae_params.ch_mult]
+    down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels)
+    up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels)
+
+    config = dict(
+        sample_size=vae_params.resolution,
+        in_channels=vae_params.in_channels,
+        out_channels=vae_params.out_ch,
+        down_block_types=tuple(down_block_types),
+        up_block_types=tuple(up_block_types),
+        block_out_channels=tuple(block_out_channels),
+        latent_channels=vae_params.z_channels,
+        layers_per_block=vae_params.num_res_blocks,
+    )
+    return config
+
+
+def create_diffusers_schedular(original_config):
+    schedular = DDIMScheduler(
+        num_train_timesteps=original_config.model.params.timesteps,
+        beta_start=original_config.model.params.linear_start,
+        beta_end=original_config.model.params.linear_end,
+        beta_schedule="scaled_linear",
+    )
+    return schedular
+
+
+def create_ldm_bert_config(original_config):
+    bert_params = original_config.model.parms.cond_stage_config.params
+    config = LDMBertConfig(
+        d_model=bert_params.n_embed,
+        encoder_layers=bert_params.n_layer,
+        encoder_ffn_dim=bert_params.n_embed * 4,
+    )
+    return config
+
+
+def convert_ldm_unet_checkpoint(checkpoint, config):
+    """
+    Takes a state dict and a config, and returns a converted checkpoint.
+    """
+
+    # extract state_dict for UNet
+    unet_state_dict = {}
+    unet_key = "model.diffusion_model."
+    keys = list(checkpoint.keys())
+    for key in keys:
+        if key.startswith(unet_key):
+            unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key)
+
+    new_checkpoint = {}
+
+    new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"]
+    new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"]
+    new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"]
+    new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"]
+
+    new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"]
+    new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"]
+
+    new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"]
+    new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"]
+    new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"]
+    new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"]
+
+    # Retrieves the keys for the input blocks only
+    num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer})
+    input_blocks = {
+        layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key]
+        for layer_id in range(num_input_blocks)
+    }
+
+    # Retrieves the keys for the middle blocks only
+    num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer})
+    middle_blocks = {
+        layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key]
+        for layer_id in range(num_middle_blocks)
+    }
+
+    # Retrieves the keys for the output blocks only
+    num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer})
+    output_blocks = {
+        layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key]
+        for layer_id in range(num_output_blocks)
+    }
+
+    for i in range(1, num_input_blocks):
+        block_id = (i - 1) // (config["layers_per_block"] + 1)
+        layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1)
+
+        resnets = [
+            key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key
+        ]
+        attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key]
+
+        if f"input_blocks.{i}.0.op.weight" in unet_state_dict:
+            new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop(
+                f"input_blocks.{i}.0.op.weight"
+            )
+            new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop(
+                f"input_blocks.{i}.0.op.bias"
+            )
+
+        paths = renew_resnet_paths(resnets)
+        meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"}
+        assign_to_checkpoint(
+            paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+        )
+
+        if len(attentions):
+            paths = renew_attention_paths(attentions)
+            meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"}
+            assign_to_checkpoint(
+                paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+            )
+
+    resnet_0 = middle_blocks[0]
+    attentions = middle_blocks[1]
+    resnet_1 = middle_blocks[2]
+
+    resnet_0_paths = renew_resnet_paths(resnet_0)
+    assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config)
+
+    resnet_1_paths = renew_resnet_paths(resnet_1)
+    assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config)
+
+    attentions_paths = renew_attention_paths(attentions)
+    meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"}
+    assign_to_checkpoint(
+        attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+    )
+
+    for i in range(num_output_blocks):
+        block_id = i // (config["layers_per_block"] + 1)
+        layer_in_block_id = i % (config["layers_per_block"] + 1)
+        output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]]
+        output_block_list = {}
+
+        for layer in output_block_layers:
+            layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1)
+            if layer_id in output_block_list:
+                output_block_list[layer_id].append(layer_name)
+            else:
+                output_block_list[layer_id] = [layer_name]
+
+        if len(output_block_list) > 1:
+            resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key]
+            attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key]
+
+            resnet_0_paths = renew_resnet_paths(resnets)
+            paths = renew_resnet_paths(resnets)
+
+            meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"}
+            assign_to_checkpoint(
+                paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+            )
+
+            if ["conv.weight", "conv.bias"] in output_block_list.values():
+                index = list(output_block_list.values()).index(["conv.weight", "conv.bias"])
+                new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[
+                    f"output_blocks.{i}.{index}.conv.weight"
+                ]
+                new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[
+                    f"output_blocks.{i}.{index}.conv.bias"
+                ]
+
+                # Clear attentions as they have been attributed above.
+                if len(attentions) == 2:
+                    attentions = []
+
+            if len(attentions):
+                paths = renew_attention_paths(attentions)
+                meta_path = {
+                    "old": f"output_blocks.{i}.1",
+                    "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}",
+                }
+                assign_to_checkpoint(
+                    paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config
+                )
+        else:
+            resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1)
+            for path in resnet_0_paths:
+                old_path = ".".join(["output_blocks", str(i), path["old"]])
+                new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]])
+
+                new_checkpoint[new_path] = unet_state_dict[old_path]
+
+    return new_checkpoint
+
+
+def convert_ldm_vae_checkpoint(checkpoint, config):
+    # extract state dict for VAE
+    vae_state_dict = {}
+    vae_key = "first_stage_model."
+    keys = list(checkpoint.keys())
+    for key in keys:
+        if key.startswith(vae_key):
+            vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key)
+
+    new_checkpoint = {}
+
+    new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]
+    new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"]
+    new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"]
+    new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"]
+    new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"]
+    new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"]
+
+    new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"]
+    new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"]
+    new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"]
+    new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"]
+    new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"]
+    new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"]
+
+    new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]
+    new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]
+    new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"]
+    new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"]
+
+    # Retrieves the keys for the encoder down blocks only
+    num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer})
+    down_blocks = {
+        layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks)
+    }
+
+    # Retrieves the keys for the decoder up blocks only
+    num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer})
+    up_blocks = {
+        layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks)
+    }
+
+    for i in range(num_down_blocks):
+        resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key]
+
+        if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:
+            new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop(
+                f"encoder.down.{i}.downsample.conv.weight"
+            )
+            new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop(
+                f"encoder.down.{i}.downsample.conv.bias"
+            )
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}
+        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+    mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]
+    num_mid_res_blocks = 2
+    for i in range(1, num_mid_res_blocks + 1):
+        resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key]
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
+        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+    mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key]
+    paths = renew_vae_attention_paths(mid_attentions)
+    meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+    conv_attn_to_linear(new_checkpoint)
+
+    for i in range(num_up_blocks):
+        block_id = num_up_blocks - 1 - i
+        resnets = [
+            key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key
+        ]
+
+        if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:
+            new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[
+                f"decoder.up.{block_id}.upsample.conv.weight"
+            ]
+            new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[
+                f"decoder.up.{block_id}.upsample.conv.bias"
+            ]
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}
+        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+    mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]
+    num_mid_res_blocks = 2
+    for i in range(1, num_mid_res_blocks + 1):
+        resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key]
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
+        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+
+    mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key]
+    paths = renew_vae_attention_paths(mid_attentions)
+    meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)
+    conv_attn_to_linear(new_checkpoint)
+    return new_checkpoint
+
+
+def convert_ldm_bert_checkpoint(checkpoint, config):
+    def _copy_attn_layer(hf_attn_layer, pt_attn_layer):
+        hf_attn_layer.q_proj.weight.data = pt_attn_layer.to_q.weight
+        hf_attn_layer.k_proj.weight.data = pt_attn_layer.to_k.weight
+        hf_attn_layer.v_proj.weight.data = pt_attn_layer.to_v.weight
+
+        hf_attn_layer.out_proj.weight = pt_attn_layer.to_out.weight
+        hf_attn_layer.out_proj.bias = pt_attn_layer.to_out.bias
+
+    def _copy_linear(hf_linear, pt_linear):
+        hf_linear.weight = pt_linear.weight
+        hf_linear.bias = pt_linear.bias
+
+    def _copy_layer(hf_layer, pt_layer):
+        # copy layer norms
+        _copy_linear(hf_layer.self_attn_layer_norm, pt_layer[0][0])
+        _copy_linear(hf_layer.final_layer_norm, pt_layer[1][0])
+
+        # copy attn
+        _copy_attn_layer(hf_layer.self_attn, pt_layer[0][1])
+
+        # copy MLP
+        pt_mlp = pt_layer[1][1]
+        _copy_linear(hf_layer.fc1, pt_mlp.net[0][0])
+        _copy_linear(hf_layer.fc2, pt_mlp.net[2])
+
+    def _copy_layers(hf_layers, pt_layers):
+        for i, hf_layer in enumerate(hf_layers):
+            if i != 0:
+                i += i
+            pt_layer = pt_layers[i : i + 2]
+            _copy_layer(hf_layer, pt_layer)
+
+    hf_model = LDMBertModel(config).eval()
+
+    # copy  embeds
+    hf_model.model.embed_tokens.weight = checkpoint.transformer.token_emb.weight
+    hf_model.model.embed_positions.weight.data = checkpoint.transformer.pos_emb.emb.weight
+
+    # copy layer norm
+    _copy_linear(hf_model.model.layer_norm, checkpoint.transformer.norm)
+
+    # copy hidden layers
+    _copy_layers(hf_model.model.layers, checkpoint.transformer.attn_layers.layers)
+
+    _copy_linear(hf_model.to_logits, checkpoint.transformer.to_logits)
+
+    return hf_model
+
+
+def convert_ldm_clip_checkpoint(checkpoint):
+    text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+
+    keys = list(checkpoint.keys())
+
+    text_model_dict = {}
+
+    for key in keys:
+        if key.startswith("cond_stage_model.transformer"):
+            text_model_dict[key[len("cond_stage_model.transformer.") :]] = checkpoint[key]
+
+    text_model.load_state_dict(text_model_dict)
+
+    return text_model
+
+import os
+def convert_checkpoint(checkpoint_path, inpainting=False):
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--checkpoint_path", default=checkpoint_path, type=str, help="Path to the checkpoint to convert."
+    )
+    # !wget https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml
+    parser.add_argument(
+        "--original_config_file",
+        default=None,
+        type=str,
+        help="The YAML config file corresponding to the original architecture.",
+    )
+    parser.add_argument(
+        "--scheduler_type",
+        default="pndm",
+        type=str,
+        help="Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim']",
+    )
+    parser.add_argument("--dump_path", default=None, type=str, help="Path to the output model.")
+
+    args = parser.parse_args([])
+    if args.original_config_file is None:
+        if inpainting:
+            args.original_config_file = "./models/v1-inpainting-inference.yaml"
+        else:
+            args.original_config_file = "./models/v1-inference.yaml"
+
+    original_config = OmegaConf.load(args.original_config_file)
+    checkpoint = torch.load(args.checkpoint_path)["state_dict"]
+
+    num_train_timesteps = original_config.model.params.timesteps
+    beta_start = original_config.model.params.linear_start
+    beta_end = original_config.model.params.linear_end
+    if args.scheduler_type == "pndm":
+        scheduler = PNDMScheduler(
+            beta_end=beta_end,
+            beta_schedule="scaled_linear",
+            beta_start=beta_start,
+            num_train_timesteps=num_train_timesteps,
+            skip_prk_steps=True,
+        )
+    elif args.scheduler_type == "lms":
+        scheduler = LMSDiscreteScheduler(beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear")
+    elif args.scheduler_type == "ddim":
+        scheduler = DDIMScheduler(
+            beta_start=beta_start,
+            beta_end=beta_end,
+            beta_schedule="scaled_linear",
+            clip_sample=False,
+            set_alpha_to_one=False,
+        )
+    else:
+        raise ValueError(f"Scheduler of type {args.scheduler_type} doesn't exist!")
+
+    # Convert the UNet2DConditionModel model.
+    unet_config = create_unet_diffusers_config(original_config)
+    converted_unet_checkpoint = convert_ldm_unet_checkpoint(checkpoint, unet_config)
+
+    unet = UNet2DConditionModel(**unet_config)
+    unet.load_state_dict(converted_unet_checkpoint)
+
+    # Convert the VAE model.
+    vae_config = create_vae_diffusers_config(original_config)
+    converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config)
+
+    vae = AutoencoderKL(**vae_config)
+    vae.load_state_dict(converted_vae_checkpoint)
+
+    # Convert the text model.
+    text_model_type = original_config.model.params.cond_stage_config.target.split(".")[-1]
+    if text_model_type == "FrozenCLIPEmbedder":
+        text_model = convert_ldm_clip_checkpoint(checkpoint)
+        tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
+        safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
+        feature_extractor = AutoFeatureExtractor.from_pretrained("CompVis/stable-diffusion-safety-checker")
+        pipe = StableDiffusionPipeline(
+            vae=vae,
+            text_encoder=text_model,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor,
+        )
+    else:
+        text_config = create_ldm_bert_config(original_config)
+        text_model = convert_ldm_bert_checkpoint(checkpoint, text_config)
+        tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
+        pipe = LDMTextToImagePipeline(vqvae=vae, bert=text_model, tokenizer=tokenizer, unet=unet, scheduler=scheduler)
+
+    return pipe

index.html

@@ -0,0 +1,404 @@
+<html>
+<head>
+<title>Stablediffusion Infinity</title>
+<meta charset="utf-8">
+<link rel="icon" type="image/x-icon" href="./favicon.png">
+
+<link rel="stylesheet" type="text/css" href="https://cdn.jsdelivr.net/gh/lkwq007/stablediffusion-infinity@v0.1.0/css/w2ui.min.css">
+<script type="text/javascript" src="https://cdn.jsdelivr.net/gh/lkwq007/stablediffusion-infinity@v0.1.0/js/w2ui.min.js"></script>
+<link rel="stylesheet" type="text/css" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.2.0/css/all.min.css">
+<script src="https://cdn.jsdelivr.net/gh/lkwq007/stablediffusion-infinity@v0.1.0/js/fabric.min.js"></script>
+<script defer src="https://cdn.jsdelivr.net/gh/lkwq007/stablediffusion-infinity@v0.1.0/js/toolbar.js"></script>
+
+<link rel="stylesheet" href="https://pyscript.net/alpha/pyscript.css" />
+<script defer src="https://pyscript.net/alpha/pyscript.js"></script>
+
+<style>
+#container {
+  position: relative;
+  margin:auto;
+  display: block;
+}
+#container > canvas {
+  position: absolute;
+  top: 0;
+  left: 0;
+}
+.control {
+  display: none;
+}
+</style>
+
+</head>
+<body>
+<div>
+<button type="button" class="control" id="export">Export</button>
+<button type="button" class="control" id="outpaint">Outpaint</button>
+<button type="button" class="control" id="undo">Undo</button>
+<button type="button" class="control" id="commit">Commit</button>
+<button type="button" class="control" id="transfer">Transfer</button>
+<button type="button" class="control" id="upload">Upload</button>
+<button type="button" class="control" id="draw">Draw</button>
+<input type="text" id="mode" value="selection" class="control">
+<input type="text" id="setup" value="0" class="control">
+<input type="text" id="upload_content" value="0" class="control">
+<textarea rows="1" id="selbuffer" name="selbuffer" class="control"></textarea>
+<fieldset class="control">
+    <div>
+      <input type="radio" id="mode0" name="mode" value="0" checked>
+      <label for="mode0">SelBox</label>
+    </div>
+    <div>
+      <input type="radio" id="mode1" name="mode" value="1">
+      <label for="mode1">Image</label>
+    </div>
+    <div>
+      <input type="radio" id="mode2" name="mode" value="2">
+      <label for="mode2">Brush</label>
+    </div>
+</fieldset>
+</div>
+<div id = "outer_container">
+<div id = "container">
+  <canvas id = "canvas0"></canvas>
+  <canvas id = "canvas1"></canvas>
+  <canvas id = "canvas2"></canvas>
+  <canvas id = "canvas3"></canvas>
+  <canvas id = "canvas4"></canvas>
+  <div id="overlay_container" style="pointer-events: none">
+    <canvas id = "overlay_canvas" width="1" height="1"></canvas>
+  </div>
+</div>
+<input type="file" name="file" id="upload_file" accept="image/*" hidden>
+<input type="file" name="state" id="upload_state" accept=".sdinf" hidden>
+<div style="position: relative;">
+<div id="toolbar" style></div>
+</div>
+</div>
+<py-env>
+- numpy
+- Pillow
+- paths:
+  - ./canvas.py
+</py-env>
+
+<py-script>
+from pyodide import to_js, create_proxy
+from PIL import Image
+import io
+import time
+import base64
+import numpy as np
+from js import (
+    console,
+    document,
+    parent,
+    devicePixelRatio,
+    ImageData,
+    Uint8ClampedArray,
+    CanvasRenderingContext2D as Context2d,
+    requestAnimationFrame,
+    window,
+    encodeURIComponent,
+    w2ui,
+    update_eraser,
+    update_scale,
+    adjust_selection,
+    update_count,
+    enable_result_lst,
+    setup_shortcut,
+)
+
+
+from canvas import InfCanvas
+
+
+
+base_lst = [None]
+async def draw_canvas() -> None:
+    width=1024
+    height=600
+    canvas=InfCanvas(1024,600)
+    update_eraser(canvas.eraser_size,min(canvas.selection_size_h,canvas.selection_size_w))
+    document.querySelector("#container").style.height= f"{height}px"
+    document.querySelector("#container").style.width = f"{width}px"
+    canvas.setup_mouse()
+    canvas.clear_background()
+    canvas.draw_buffer()
+    canvas.draw_selection_box()
+    base_lst[0]=canvas
+
+async def draw_canvas_func():
+
+    width=1500
+    height=600
+    selection_size=256
+    document.querySelector("#container").style.width = f"{width}px"
+    document.querySelector("#container").style.height= f"{height}px"
+    canvas=InfCanvas(int(width),int(height),selection_size=int(selection_size))
+    canvas.setup_mouse()
+    canvas.clear_background()
+    canvas.draw_buffer()
+    canvas.draw_selection_box()
+    base_lst[0]=canvas
+  
+async def export_func(event):
+    base=base_lst[0]
+    arr=base.export()
+    base.draw_buffer()
+    base.canvas[2].clear()
+    base64_str = base.numpy_to_base64(arr)
+    time_str = time.strftime("%Y%m%d_%H%M%S")
+    link = document.createElement("a")
+    if len(event.data)>2 and event.data[2]:
+        filename = event.data[2]
+    else:
+        filename = f"outpaint_{time_str}"
+    # link.download = f"sdinf_state_{time_str}.json"
+    link.download = f"{filename}.png"
+    # link.download = f"outpaint_{time_str}.png"
+    link.href = "data:image/png;base64,"+base64_str
+    link.click()
+    console.log(f"Canvas saved to {filename}.png")
+
+img_candidate_lst=[None,0]
+
+async def outpaint_func(event):
+    base=base_lst[0]
+    if len(event.data)==2:
+        app=parent.document.querySelector("gradio-app")
+        if app.shadowRoot:
+            app=app.shadowRoot
+        base64_str_raw=app.querySelector("#output textarea").value
+        base64_str_lst=base64_str_raw.split(",")
+        img_candidate_lst[0]=base64_str_lst
+        img_candidate_lst[1]=0
+    elif event.data[2]=="next":
+        img_candidate_lst[1]+=1
+    elif event.data[2]=="prev":
+        img_candidate_lst[1]-=1
+    enable_result_lst()
+    if img_candidate_lst[0] is None:
+        return
+    lst=img_candidate_lst[0]
+    idx=img_candidate_lst[1]
+    update_count(idx%len(lst)+1,len(lst))
+    arr=base.base64_to_numpy(lst[idx%len(lst)])
+    base.fill_selection(arr)
+    base.draw_selection_box()
+
+async def undo_func(event):
+    base=base_lst[0]
+    img_candidate_lst[0]=None
+    if base.sel_dirty:
+        base.sel_buffer = np.zeros((base.selection_size_h, base.selection_size_w, 4), dtype=np.uint8)
+        base.sel_dirty = False
+    base.canvas[2].clear()
+
+async def commit_func(event):
+    base=base_lst[0]
+    img_candidate_lst[0]=None
+    if base.sel_dirty:
+        base.write_selection_to_buffer()
+        base.draw_buffer()
+    base.canvas[2].clear()
+
+async def transfer_func(event):
+    base=base_lst[0]
+    base.read_selection_from_buffer()
+    sel_buffer=base.sel_buffer
+    sel_buffer_str=base.numpy_to_base64(sel_buffer)
+    app=parent.document.querySelector("gradio-app")
+    if app.shadowRoot:
+        app=app.shadowRoot
+    app.querySelector("#input textarea").value=sel_buffer_str
+    app.querySelector("#proceed").click()
+
+async def upload_func(event):
+    base=base_lst[0]
+    # base64_str=event.data[1]
+    base64_str=document.querySelector("#upload_content").value
+    base64_str=base64_str.split(",")[-1]
+    # base64_str=parent.document.querySelector("gradio-app").shadowRoot.querySelector("#upload textarea").value
+    arr=base.base64_to_numpy(base64_str)
+    h,w,c=base.buffer.shape
+    base.sync_to_buffer()
+    base.buffer_dirty=True
+    mask=arr[:,:,3:4].repeat(4,axis=2)
+    base.buffer[mask>0]=0
+    # in case mismatch
+    base.buffer[0:h,0:w,:]+=arr
+    #base.buffer[yo:yo+h,xo:xo+w,0:3]=arr[:,:,0:3]
+    #base.buffer[yo:yo+h,xo:xo+w,-1]=arr[:,:,-1]
+    base.draw_buffer()
+
+async def setup_shortcut_func(event):
+    setup_shortcut(event.data[1])
+  
+
+document.querySelector("#export").addEventListener("click",create_proxy(export_func))
+document.querySelector("#undo").addEventListener("click",create_proxy(undo_func))
+document.querySelector("#commit").addEventListener("click",create_proxy(commit_func))
+document.querySelector("#outpaint").addEventListener("click",create_proxy(outpaint_func))
+document.querySelector("#upload").addEventListener("click",create_proxy(upload_func))
+
+document.querySelector("#transfer").addEventListener("click",create_proxy(transfer_func))
+document.querySelector("#draw").addEventListener("click",create_proxy(draw_canvas_func))
+
+async def setup_func():
+    document.querySelector("#setup").value="1"
+
+async def reset_func(event):
+    base=base_lst[0]
+    base.reset()
+    
+async def load_func(event):
+    base=base_lst[0]
+    base.load(event.data[1])
+
+async def save_func(event):
+    base=base_lst[0]
+    json_str=base.save()
+    time_str = time.strftime("%Y%m%d_%H%M%S")
+    link = document.createElement("a")
+    if len(event.data)>2 and event.data[2]:
+        filename = str(event.data[2]).strip()
+    else:
+        filename = f"outpaint_{time_str}"
+    # link.download = f"sdinf_state_{time_str}.json"
+    link.download = f"{filename}.sdinf"
+    link.href = "data:text/json;charset=utf-8,"+encodeURIComponent(json_str)
+    link.click()
+
+async def prev_result_func(event):
+    base=base_lst[0]
+    base.reset()
+
+async def next_result_func(event):
+    base=base_lst[0]
+    base.reset()
+
+async def zoom_in_func(event):
+    base=base_lst[0]
+    scale=base.scale
+    if scale>=0.2:
+        scale-=0.1
+        if len(event.data)>2:
+            base.update_scale(scale,int(event.data[2]),int(event.data[3]))
+        else:
+            base.update_scale(scale)
+        scale=base.scale
+        update_scale(f"{base.width}x{base.height} ({round(100/scale)}%)")
+
+async def zoom_out_func(event):
+    base=base_lst[0]
+    scale=base.scale
+    if scale<10:
+        scale+=0.1
+        console.log(len(event.data))
+        if len(event.data)>2:
+            base.update_scale(scale,int(event.data[2]),int(event.data[3]))
+        else:
+            base.update_scale(scale)
+        scale=base.scale
+        update_scale(f"{base.width}x{base.height} ({round(100/scale)}%)")
+
+async def sync_func(event):
+    base=base_lst[0]
+    base.sync_to_buffer()
+    base.canvas[2].clear()
+
+async def eraser_size_func(event):
+    base=base_lst[0]
+    eraser_size=min(int(event.data[1]),min(base.selection_size_h,base.selection_size_w))
+    eraser_size=max(8,eraser_size)
+    base.eraser_size=eraser_size
+
+async def resize_selection_func(event):
+    base=base_lst[0]
+    cursor=base.cursor
+    if len(event.data)>3:
+        console.log(event.data)
+        base.cursor[0]=int(event.data[1])
+        base.cursor[1]=int(event.data[2])
+        base.selection_size_w=int(event.data[3])//8*8
+        base.selection_size_h=int(event.data[4])//8*8
+        base.refine_selection()
+        base.draw_selection_box()
+    elif len(event.data)>2:
+        base.draw_selection_box()
+    else:
+        base.canvas[-1].clear()
+        adjust_selection(cursor[0],cursor[1],base.selection_size_w,base.selection_size_h)
+
+async def eraser_func(event):
+    base=base_lst[0]
+    if event.data[1]!="eraser":
+        base.canvas[-2].clear()
+    else:
+        x,y=base.mouse_pos
+        base.draw_eraser(x,y)
+
+async def resize_func(event):
+    base=base_lst[0]
+    width=int(event.data[1])
+    height=int(event.data[2])
+    if width>=256 and height>=256:
+        if max(base.selection_size_h,base.selection_size_w)>min(width,height):
+            base.selection_size_h=256
+            base.selection_size_w=256
+        base.resize(width,height)
+
+async def message_func(event):
+    if event.data[0]=="click":
+        if event.data[1]=="clear":
+            await reset_func(event)
+        elif event.data[1]=="save":
+            await save_func(event)
+        elif event.data[1]=="export":
+            await export_func(event)
+        elif event.data[1]=="accept":
+            await commit_func(event)
+        elif event.data[1]=="cancel":
+            await undo_func(event)
+        elif event.data[1]=="zoom_in":
+            await zoom_in_func(event)
+        elif event.data[1]=="zoom_out":
+            await zoom_out_func(event)
+    elif event.data[0]=="sync":
+        await sync_func(event)
+    elif event.data[0]=="load":
+        await load_func(event)
+    elif event.data[0]=="upload":
+        await upload_func(event)
+    elif event.data[0]=="outpaint":
+        await outpaint_func(event)
+    elif event.data[0]=="mode":
+        if event.data[1]!="selection":
+            await sync_func(event)
+        await eraser_func(event)
+        document.querySelector("#mode").value=event.data[1]
+    elif event.data[0]=="transfer":
+        await transfer_func(event)
+    elif event.data[0]=="setup":
+        await draw_canvas_func(event)
+    elif event.data[0]=="eraser_size":
+        await eraser_size_func(event)
+    elif event.data[0]=="resize_selection":
+        await resize_selection_func(event)
+    elif event.data[0]=="shortcut":
+        await setup_shortcut_func(event)
+    elif event.data[0]=="resize":
+        await resize_func(event)
+    
+window.addEventListener("message",create_proxy(message_func))
+
+import asyncio
+
+_ = await asyncio.gather(
+  setup_func(),draw_canvas_func()
+)
+</py-script>
+
+</body>
+</html>

js/keyboard.js

@@ -0,0 +1,37 @@
+
+window.my_setup_keyboard=setInterval(function(){
+    let app=document.querySelector("gradio-app");
+    app=app.shadowRoot??app;
+    let frame=app.querySelector("#sdinfframe").contentWindow;
+    console.log("Check iframe...");
+    if(frame.setup_shortcut)
+    {
+        frame.setup_shortcut(json);
+        clearInterval(window.my_setup_keyboard);
+    }
+}, 1000);
+var config=JSON.parse(json);
+var key_map={};
+Object.keys(config.shortcut).forEach(k=>{
+    key_map[config.shortcut[k]]=k;
+});
+document.addEventListener("keydown", e => {
+    if(e.target.tagName!="INPUT"&&e.target.tagName!="GRADIO-APP"&&e.target.tagName!="TEXTAREA")
+    {
+        let key=e.key;
+        if(e.ctrlKey)
+        {
+            key="Ctrl+"+e.key;
+            if(key in key_map)
+            {
+                e.preventDefault();
+            }
+        }
+        let app=document.querySelector("gradio-app");
+        app=app.shadowRoot??app;
+        let frame=app.querySelector("#sdinfframe").contentDocument;
+        frame.dispatchEvent(
+          new KeyboardEvent("keydown", {key: e.key, ctrlKey: e.ctrlKey})
+        );
+    }
+})

js/mode.js

@@ -0,0 +1,6 @@
+function(mode){
+    let app=document.querySelector("gradio-app").shadowRoot;
+    let frame=app.querySelector("#sdinfframe").contentWindow.document;
+    frame.querySelector("#mode").value=mode;
+    return mode;
+}

js/outpaint.js

@@ -0,0 +1,23 @@
+function(a){
+    if(!window.my_observe_outpaint)
+    {
+        console.log("setup outpaint here");
+        window.my_observe_outpaint = new MutationObserver(function (event) {
+            console.log(event);
+            let app=document.querySelector("gradio-app");
+            app=app.shadowRoot??app;
+            let frame=app.querySelector("#sdinfframe").contentWindow;
+            frame.postMessage(["outpaint", ""], "*");
+        });
+        var app=document.querySelector("gradio-app");
+        app=app.shadowRoot??app;
+        window.my_observe_outpaint_target=app.querySelector("#output span");
+        window.my_observe_outpaint.observe(window.my_observe_outpaint_target, {
+            attributes: false, 
+            subtree: true,
+            childList: true, 
+            characterData: true
+        });
+    }
+    return a;
+}

js/proceed.js

@@ -0,0 +1,42 @@
+function(sel_buffer_str,
+    prompt_text,
+    negative_prompt_text,
+    strength,
+    guidance,
+    step,
+    resize_check,
+    fill_mode,
+    enable_safety,
+    use_correction,
+    enable_img2img,
+    use_seed,
+    seed_val,
+    generate_num,
+    scheduler,
+    scheduler_eta,
+    state){
+    let app=document.querySelector("gradio-app");
+    app=app.shadowRoot??app;
+    sel_buffer=app.querySelector("#input textarea").value;
+    let use_correction_bak=false;
+    ({resize_check,enable_safety,use_correction_bak,enable_img2img,use_seed,seed_val}=window.config_obj);
+    return [
+        sel_buffer,
+        prompt_text,
+        negative_prompt_text,
+        strength,
+        guidance,
+        step,
+        resize_check,
+        fill_mode,
+        enable_safety,
+        use_correction,
+        enable_img2img,
+        use_seed,
+        seed_val,
+        generate_num,
+        scheduler,
+        scheduler_eta,
+        state,
+    ]
+}

js/setup.js

@@ -0,0 +1,28 @@
+function(token_val, width, height, size, model_choice, model_path){
+    let app=document.querySelector("gradio-app");
+    app=app.shadowRoot??app;
+    app.querySelector("#sdinfframe").style.height=80+Number(height)+"px";
+    // app.querySelector("#setup_row").style.display="none";
+    app.querySelector("#model_path_input").style.display="none";
+    let frame=app.querySelector("#sdinfframe").contentWindow.document;
+
+    if(frame.querySelector("#setup").value=="0")
+    {
+        window.my_setup=setInterval(function(){
+            let app=document.querySelector("gradio-app");
+            app=app.shadowRoot??app;
+            let frame=app.querySelector("#sdinfframe").contentWindow.document;
+            console.log("Check PyScript...")
+            if(frame.querySelector("#setup").value=="1")
+            {
+                frame.querySelector("#draw").click();
+                clearInterval(window.my_setup);
+            }
+        }, 100)
+    }
+    else
+    {
+        frame.querySelector("#draw").click();
+    }
+    return [token_val, width, height, size, model_choice, model_path];
+}

js/toolbar.js

@@ -0,0 +1,581 @@
+// import { w2ui,w2toolbar,w2field,query,w2alert, w2utils,w2confirm} from "https://rawgit.com/vitmalina/w2ui/master/dist/w2ui.es6.min.js"
+// import { w2ui,w2toolbar,w2field,query,w2alert, w2utils,w2confirm} from "https://cdn.jsdelivr.net/gh/vitmalina/w2ui@master/dist/w2ui.es6.min.js"
+
+// https://stackoverflow.com/questions/36280818/how-to-convert-file-to-base64-in-javascript
+function getBase64(file) {
+   var reader = new FileReader();
+   reader.readAsDataURL(file);
+   reader.onload = function () {
+    add_image(reader.result);
+    //  console.log(reader.result);
+   };
+   reader.onerror = function (error) {
+     console.log("Error: ", error);
+   };
+}
+
+function getText(file) {
+   var reader = new FileReader();
+   reader.readAsText(file);
+   reader.onload = function () {
+    window.postMessage(["load",reader.result],"*")
+    //  console.log(reader.result);
+   };
+   reader.onerror = function (error) {
+     console.log("Error: ", error);
+   };
+}
+
+document.querySelector("#upload_file").addEventListener("change", (event)=>{
+    console.log(event);
+    let file = document.querySelector("#upload_file").files[0];
+    getBase64(file);
+})
+
+document.querySelector("#upload_state").addEventListener("change", (event)=>{
+    console.log(event);
+    let file = document.querySelector("#upload_state").files[0];
+    getText(file);
+})
+
+open_setting = function() {
+    if (!w2ui.foo) {
+        new w2form({
+            name: "foo",
+            style: "border: 0px; background-color: transparent;",
+            fields: [{
+                    field: "canvas_width",
+                    type: "int",
+                    required: true,
+                    html: {
+                        label: "Canvas Width"
+                    }
+                },
+                {
+                    field: "canvas_height",
+                    type: "int",
+                    required: true,
+                    html: {
+                        label: "Canvas Height"
+                    }
+                },
+            ],
+            record: {
+                canvas_width: 1200,
+                canvas_height: 600,
+            },
+            actions: {
+                Save() {
+                    this.validate();
+                    let record = this.getCleanRecord();
+                    window.postMessage(["resize",record.canvas_width,record.canvas_height],"*");
+                    w2popup.close();
+                },
+                custom: {
+                    text: "Cancel",
+                    style: "text-transform: uppercase",
+                    onClick(event) {
+                        w2popup.close();
+                    }
+                }
+            }
+        });
+    }
+    w2popup.open({
+            title: "Form in a Popup",
+            body: "<div id='form' style='width: 100%; height: 100%;''></div>",
+            style: "padding: 15px 0px 0px 0px",
+            width: 500,
+            height: 280,
+            showMax: true,
+            async onToggle(event) {
+                await event.complete
+                w2ui.foo.resize();
+            }
+        })
+        .then((event) => {
+            w2ui.foo.render("#form")
+        });
+}
+
+var button_lst=["clear", "load", "save", "export", "upload", "selection", "canvas", "eraser", "outpaint", "accept", "cancel", "retry", "prev", "current", "next", "eraser_size_btn", "eraser_size", "resize_selection", "scale", "zoom_in", "zoom_out", "help"];
+var upload_button_lst=['clear', 'load', 'save', "upload", 'export', 'outpaint', 'resize_selection', 'help', "setting"];
+var resize_button_lst=['clear', 'load', 'save', "upload", 'export', "selection", "canvas", "eraser", 'outpaint', 'resize_selection',"zoom_in", "zoom_out", 'help', "setting"];
+var outpaint_button_lst=['clear', 'load', 'save', "canvas", "eraser", "upload", 'export', 'resize_selection', "zoom_in", "zoom_out",'help', "setting"];
+var outpaint_result_lst=["accept", "cancel", "retry", "prev", "current", "next"];
+var outpaint_result_func_lst=["accept", "retry", "prev", "current", "next"];
+
+function check_button(id,text="",checked=true,tooltip="")
+{
+    return { type: "check",  id: id, text: text, icon: checked?"fa-solid fa-square-check":"fa-regular fa-square", checked: checked, tooltip: tooltip };
+}
+
+var toolbar=new w2toolbar({
+    box: "#toolbar",
+    name: "toolbar",
+    tooltip: "top",
+    items: [
+        { type: "button", id: "clear", text: "Reset", tooltip: "Reset Canvas", icon: "fa-solid fa-rectangle-xmark" },
+        { type: "break" },
+        { type: "button", id: "load", tooltip: "Load Canvas", icon: "fa-solid fa-file-import" },
+        { type: "button", id: "save", tooltip: "Save Canvas", icon: "fa-solid fa-file-export" },
+        { type: "button", id: "export", tooltip: "Export Image", icon: "fa-solid fa-floppy-disk" },
+        { type: "break" },
+        { type: "button", id: "upload", text: "Upload Image", icon: "fa-solid fa-upload" },
+        { type: "break" },
+        { type: "radio", id: "selection", group: "1", tooltip: "Selection", icon: "fa-solid fa-arrows-up-down-left-right", checked: true },
+        { type: "radio", id: "canvas", group: "1", tooltip: "Canvas", icon: "fa-solid fa-image" },
+        { type: "radio", id: "eraser", group: "1", tooltip: "Eraser", icon: "fa-solid fa-eraser" },
+        { type: "break" },
+        { type: "button", id: "outpaint", text: "Outpaint", tooltip: "Run Outpainting", icon: "fa-solid fa-brush" },
+        { type: "break" },
+        { type: "button", id: "accept", text: "Accept", tooltip: "Accept current result", icon: "fa-solid fa-check", hidden: true, disable:true,},
+        { type: "button", id: "cancel", text: "Cancel", tooltip: "Cancel current outpainting/error", icon: "fa-solid fa-ban", hidden: true},
+        { type: "button", id: "retry", text: "Retry", tooltip: "Retry", icon: "fa-solid fa-rotate", hidden: true, disable:true,},
+        { type: "button", id: "prev", tooltip: "Prev Result", icon: "fa-solid fa-caret-left", hidden: true, disable:true,},
+        { type: "html", id: "current", hidden: true, disable:true,
+            async onRefresh(event) {
+                await event.complete
+                let fragment = query.html(`
+                <div class="w2ui-tb-text">
+                <div class="w2ui-tb-count">
+                    <span>${this.sel_value ?? "1/1"}</span>
+                </div> </div>`)
+                query(this.box).find("#tb_toolbar_item_current").append(fragment)
+            }
+        },
+        { type: "button", id: "next", tooltip: "Next Result", icon: "fa-solid fa-caret-right", hidden: true,disable:true,},
+        { type: "button", id: "add_image", text: "Add Image", icon: "fa-solid fa-file-circle-plus", hidden: true,disable:true,},
+        { type: "button", id: "delete_image", text: "Delete Image", icon: "fa-solid fa-trash-can", hidden: true,disable:true,},
+        { type: "button", id: "confirm", text: "Confirm", icon: "fa-solid fa-check", hidden: true,disable:true,},
+        { type: "button", id: "cancel_overlay", text: "Cancel", icon: "fa-solid fa-ban", hidden: true,disable:true,},
+        { type: "break" },
+        { type: "spacer" },
+        { type: "break" },
+        { type: "button", id: "eraser_size_btn", tooltip: "Eraser Size", text:"Size", icon: "fa-solid fa-eraser", hidden: true, count: 32},
+        { type: "html", id: "eraser_size", hidden: true,
+            async onRefresh(event) {
+                await event.complete
+                // let fragment = query.html(`
+                //     <input type="number" size="${this.eraser_size ? this.eraser_size.length:"2"}" style="margin: 0px 3px; padding: 4px;" min="8" max="${this.eraser_max ?? "256"}" value="${this.eraser_size ?? "32"}">
+                //     <input type="range" style="margin: 0px 3px; padding: 4px;" min="8" max="${this.eraser_max ?? "256"}" value="${this.eraser_size ?? "32"}">`)
+                let fragment = query.html(`
+                    <input type="range" style="margin: 0px 3px; padding: 4px;" min="8" max="${this.eraser_max ?? "256"}" value="${this.eraser_size ?? "32"}">
+                    `)
+                fragment.filter("input").on("change", event => {
+                    this.eraser_size = event.target.value;
+                    window.overlay.freeDrawingBrush.width=this.eraser_size;
+                    this.setCount("eraser_size_btn", event.target.value);
+                    window.postMessage(["eraser_size", event.target.value],"*")
+                    this.refresh();
+                })
+                query(this.box).find("#tb_toolbar_item_eraser_size").append(fragment)
+            }
+        },
+        // { type: "button", id: "resize_eraser", tooltip: "Resize Eraser", icon: "fa-solid fa-sliders" },
+        { type: "button", id: "resize_selection", text: "Resize Selection", tooltip: "Resize Selection", icon: "fa-solid fa-expand" },
+        { type: "break" },
+        { type: "html", id: "scale",
+            async onRefresh(event) {
+                await event.complete
+                let fragment = query.html(`
+                <div class="">
+                <div style="padding: 4px; border: 1px solid silver">
+                    <span>${this.scale_value ?? "100%"}</span>
+                </div></div>`)
+                query(this.box).find("#tb_toolbar_item_scale").append(fragment)
+            }
+        },
+        { type: "button", id: "zoom_in", tooltip: "Zoom In", icon: "fa-solid fa-magnifying-glass-plus" },
+        { type: "button", id: "zoom_out", tooltip: "Zoom Out", icon: "fa-solid fa-magnifying-glass-minus" },
+        { type: "break" },
+        { type: "button", id: "help", tooltip: "Help", icon: "fa-solid fa-circle-info" },
+        { type: "new-line"},
+        { type: "button", id: "setting", text: "Canvas Setting", tooltip: "Resize Canvas Here", icon: "fa-solid fa-sliders" },
+        { type: "break" },
+        check_button("enable_img2img","Enable Img2Img",false),
+        // check_button("use_correction","Photometric Correction",false),
+        check_button("resize_check","Resize Small Input",true),
+        check_button("enable_safety","Enable Safety Checker",true),
+        check_button("square_selection","Square Selection Only",false),
+        {type: "break"},
+        check_button("use_seed","Use Seed:",false),
+        { type: "html", id: "seed_val",
+            async onRefresh(event) {
+                await event.complete
+                let fragment = query.html(`
+                    <input type="number" style="margin: 0px 3px; padding: 4px; width:100px;" value="${this.config_obj.seed_val ?? "0"}">`)
+                fragment.filter("input").on("change", event => {
+                    this.config_obj.seed_val = event.target.value;
+                    parent.config_obj=this.config_obj;
+                    this.refresh();
+                })
+                query(this.box).find("#tb_toolbar_item_seed_val").append(fragment)
+            }
+        },
+        { type: "button", id: "random_seed", tooltip: "Set a random seed", icon: "fa-solid fa-dice" },
+    ],
+    onClick(event) {
+        switch(event.target){
+            case "setting":
+                open_setting();
+                break;
+            case "upload":
+                this.upload_mode=true
+                document.querySelector("#overlay_container").style.pointerEvents="auto";
+                this.click("canvas");
+                this.click("selection");
+                this.show("confirm","cancel_overlay","add_image","delete_image");
+                this.enable("confirm","cancel_overlay","add_image","delete_image");
+                this.disable(...upload_button_lst);
+                query("#upload_file").click();
+                if(this.upload_tip)
+                {
+                    this.upload_tip=false;
+                    w2utils.notify("Note that only visible images will be added to canvas",{timeout:10000,where:query("#container")})
+                }
+                break;
+            case "resize_selection":
+                this.resize_mode=true;
+                this.disable(...resize_button_lst);
+                this.enable("confirm","cancel_overlay");
+                this.show("confirm","cancel_overlay");
+                window.postMessage(["resize_selection",""],"*");
+                document.querySelector("#overlay_container").style.pointerEvents="auto";
+                break;
+            case "confirm":
+                if(this.upload_mode)
+                {
+                    export_image();
+                }
+                else
+                {
+                    let sel_box=this.selection_box;
+                    window.postMessage(["resize_selection",sel_box.x,sel_box.y,sel_box.width,sel_box.height],"*");
+                }
+            case "cancel_overlay":
+                end_overlay();
+                this.hide("confirm","cancel_overlay","add_image","delete_image");
+                if(this.upload_mode){
+                    this.enable(...upload_button_lst);
+                }
+                else
+                {
+                    this.enable(...resize_button_lst);
+                    window.postMessage(["resize_selection","",""],"*");
+                    if(event.target=="cancel_overlay")
+                    {
+                        this.selection_box=this.selection_box_bak;
+                    }
+                }
+                if(this.selection_box)
+                {
+                    this.setCount("resize_selection",`${Math.floor(this.selection_box.width/8)*8}x${Math.floor(this.selection_box.height/8)*8}`);
+                }
+                this.disable("confirm","cancel_overlay","add_image","delete_image");
+                this.upload_mode=false;
+                this.resize_mode=false;
+                this.click("selection");
+                break;
+            case "add_image":
+                query("#upload_file").click();
+                break;
+            case "delete_image":
+                let active_obj = window.overlay.getActiveObject();
+                if(active_obj)
+                {
+                    window.overlay.remove(active_obj);
+                    window.overlay.renderAll();
+                }
+                else
+                {
+                    w2utils.notify("You need to select an image first",{error:true,timeout:2000,where:query("#container")})
+                }
+                break;
+            case "load":
+                query("#upload_state").click();
+                this.selection_box=null;
+                this.setCount("resize_selection","");
+                break;
+            case "next":
+            case "prev":
+                window.postMessage(["outpaint", "", event.target], "*");
+                break;
+            case "outpaint":
+                this.click("selection");
+                this.disable(...outpaint_button_lst);
+                this.show(...outpaint_result_lst);
+                if(this.outpaint_tip)
+                {
+                    this.outpaint_tip=false;
+                    w2utils.notify("The canvas stays locked until you accept/cancel current outpainting",{timeout:10000,where:query("#container")})
+                }
+                document.querySelector("#container").style.pointerEvents="none";
+            case "retry":
+                this.disable(...outpaint_result_func_lst);
+                window.postMessage(["transfer",""],"*")
+                break;
+            case "accept":
+            case "cancel":
+                this.hide(...outpaint_result_lst);
+                this.disable(...outpaint_result_func_lst);
+                this.enable(...outpaint_button_lst);
+                document.querySelector("#container").style.pointerEvents="auto";
+                window.postMessage(["click", event.target],"*");
+                let app=parent.document.querySelector("gradio-app");
+                app=app.shadowRoot??app;
+                app.querySelector("#cancel").click();
+                break;
+            case "eraser":
+            case "selection":
+            case "canvas":
+                if(event.target=="eraser")
+                {
+                    this.show("eraser_size","eraser_size_btn");
+                    window.overlay.freeDrawingBrush.width=this.eraser_size;
+                    window.overlay.isDrawingMode = true;
+                }
+                else
+                {
+                    this.hide("eraser_size","eraser_size_btn");
+                    window.overlay.isDrawingMode = false;
+                }
+                if(this.upload_mode)
+                {
+                    if(event.target=="canvas")
+                    {
+                        window.postMessage(["mode", event.target],"*")
+                        document.querySelector("#overlay_container").style.pointerEvents="none";
+                        document.querySelector("#overlay_container").style.opacity = 0.5;
+                    }
+                    else
+                    {
+                        document.querySelector("#overlay_container").style.pointerEvents="auto";
+                        document.querySelector("#overlay_container").style.opacity = 1.0;
+                    }
+                }
+                else
+                {
+                    window.postMessage(["mode", event.target],"*")
+                }
+                break;
+            case "help":
+                w2popup.open({
+                    title: "Document",
+                    body: "Usage: <a href='https://github.com/lkwq007/stablediffusion-infinity/blob/master/docs/usage.md'  target='_blank'>https://github.com/lkwq007/stablediffusion-infinity/blob/master/docs/usage.md</a>"
+                })
+                break;
+            case "clear":
+                w2confirm("Reset canvas?").yes(() => {
+                    window.postMessage(["click", event.target],"*");
+                }).no(() => {})
+                break;
+            case "random_seed":
+                this.config_obj.seed_val=Math.floor(Math.random() * 3000000000);
+                parent.config_obj=this.config_obj;
+                this.refresh();
+                break;
+            case "enable_img2img":
+            case "use_correction":
+            case "resize_check":
+            case "enable_safety":
+            case "use_seed":
+            case "square_selection":
+                let target=this.get(event.target);
+                target.icon=target.checked?"fa-regular fa-square":"fa-solid fa-square-check";
+                this.config_obj[event.target]=!target.checked;
+                parent.config_obj=this.config_obj;
+                this.refresh();
+                break;
+            case "save":
+            case "export":
+                ask_filename(event.target);
+                break;
+            default:
+                // clear, save, export, outpaint, retry
+                // break, save, export, accept, retry, outpaint
+                window.postMessage(["click", event.target],"*")
+        }
+        console.log("Target: "+ event.target, event)
+    }
+})
+window.w2ui=w2ui;
+w2ui.toolbar.config_obj={
+    resize_check: true,
+    enable_safety: true,
+    use_correction: false,
+    enable_img2img: false,
+    use_seed: false,
+    seed_val: 0,
+    square_selection: false,
+};
+w2ui.toolbar.outpaint_tip=true;
+w2ui.toolbar.upload_tip=true;
+window.update_count=function(cur,total){
+  w2ui.toolbar.sel_value=`${cur}/${total}`;
+  w2ui.toolbar.refresh();
+}
+window.update_eraser=function(val,max_val){
+  w2ui.toolbar.eraser_size=`${val}`;
+  w2ui.toolbar.eraser_max=`${max_val}`;
+  w2ui.toolbar.setCount("eraser_size_btn", `${val}`);
+  w2ui.toolbar.refresh();
+}
+window.update_scale=function(val){
+  w2ui.toolbar.scale_value=`${val}`;
+  w2ui.toolbar.refresh();
+}
+window.enable_result_lst=function(){
+  w2ui.toolbar.enable(...outpaint_result_lst);
+}
+function onObjectScaled(e)
+{
+    let object = e.target;
+    if(object.isType("rect"))
+    {
+        let width=object.getScaledWidth();
+        let height=object.getScaledHeight();
+        object.scale(1);
+        width=Math.max(Math.min(width,window.overlay.width-object.left),256);
+        height=Math.max(Math.min(height,window.overlay.height-object.top),256);
+        let l=Math.max(Math.min(object.left,window.overlay.width-width-object.strokeWidth),0);
+        let t=Math.max(Math.min(object.top,window.overlay.height-height-object.strokeWidth),0);
+        if(window.w2ui.toolbar.config_obj.square_selection)
+        {
+            let max_val = Math.min(Math.max(width,height),window.overlay.width,window.overlay.height);
+            width=max_val;
+            height=max_val;
+        }
+        object.set({ width: width, height: height, left:l,top:t})
+        window.w2ui.toolbar.selection_box={width: width, height: height, x:object.left, y:object.top};
+        window.w2ui.toolbar.setCount("resize_selection",`${Math.floor(width/8)*8}x${Math.floor(height/8)*8}`);
+        window.w2ui.toolbar.refresh();
+    }
+}
+function onObjectMoved(e)
+{
+    let object = e.target;
+    if(object.isType("rect"))
+    {   
+        let l=Math.max(Math.min(object.left,window.overlay.width-object.width-object.strokeWidth),0);
+        let t=Math.max(Math.min(object.top,window.overlay.height-object.height-object.strokeWidth),0);
+        object.set({left:l,top:t});
+        window.w2ui.toolbar.selection_box={width: object.width, height: object.height, x:object.left, y:object.top};
+    }
+}
+window.setup_overlay=function(width,height)
+{
+    if(window.overlay)
+    {
+        window.overlay.setDimensions({width:width,height:height});
+        let app=parent.document.querySelector("gradio-app");
+        app=app.shadowRoot??app;
+        app.querySelector("#sdinfframe").style.height=80+Number(height)+"px";
+        document.querySelector("#container").style.height= height+"px";
+        document.querySelector("#container").style.width = width+"px";
+    }
+    else
+    {
+        canvas=new fabric.Canvas("overlay_canvas");
+        canvas.setDimensions({width:width,height:height});
+        let app=parent.document.querySelector("gradio-app");
+        app=app.shadowRoot??app;
+        app.querySelector("#sdinfframe").style.height=80+Number(height)+"px";
+        canvas.freeDrawingBrush = new fabric.EraserBrush(canvas);
+        canvas.on("object:scaling", onObjectScaled);
+        canvas.on("object:moving", onObjectMoved);
+        window.overlay=canvas;
+    }
+    document.querySelector("#overlay_container").style.pointerEvents="none";
+}
+window.update_overlay=function(width,height)
+{
+    window.overlay.setDimensions({width:width,height:height},{backstoreOnly:true});
+    // document.querySelector("#overlay_container").style.pointerEvents="none";
+}
+window.adjust_selection=function(x,y,width,height)
+{
+    var rect = new fabric.Rect({
+        left: x,
+        top: y,
+        fill: "rgba(0,0,0,0)",
+        strokeWidth: 3, 
+        stroke: "rgba(0,0,0,0.7)",
+        cornerColor: "red",
+        cornerStrokeColor: "red",
+        borderColor: "rgba(255, 0, 0, 1.0)",
+        width: width,
+        height: height,
+        lockRotation: true,
+    });
+    rect.setControlsVisibility({ mtr: false });
+    window.overlay.add(rect);
+    window.overlay.setActiveObject(window.overlay.item(0));
+    window.w2ui.toolbar.selection_box={width: width, height: height, x:x, y:y};
+    window.w2ui.toolbar.selection_box_bak={width: width, height: height, x:x, y:y};
+}
+function add_image(url)
+{
+    fabric.Image.fromURL(url,function(img){
+        window.overlay.add(img);
+        window.overlay.setActiveObject(img);
+    },{left:100,top:100});
+}
+function export_image()
+{
+    data=window.overlay.toDataURL();
+    document.querySelector("#upload_content").value=data;
+    window.postMessage(["upload",""],"*");
+    end_overlay();
+}
+function end_overlay()
+{
+    window.overlay.clear();
+    document.querySelector("#overlay_container").style.opacity = 1.0;
+    document.querySelector("#overlay_container").style.pointerEvents="none";
+}
+function ask_filename(target)
+{
+    w2prompt({
+        label: "Enter filename",
+        value: `outpaint_${((new Date(Date.now() -(new Date()).getTimezoneOffset() * 60000))).toISOString().replace("T","_").replace(/[^0-9_]/g, "").substring(0,15)}`,
+    })
+    .change((event) => {
+        console.log("change", event.detail.originalEvent.target.value);
+    })
+    .ok((event) => {
+        console.log("value=", event.detail.value);
+        window.postMessage(["click",target,event.detail.value],"*");
+    })
+    .cancel((event) => {
+        console.log("cancel");
+    });
+}
+
+document.querySelector("#container").addEventListener("wheel",(e)=>{e.preventDefault()})
+window.setup_shortcut=function(json)
+{
+    var config=JSON.parse(json);
+    var key_map={};
+    Object.keys(config.shortcut).forEach(k=>{
+        key_map[config.shortcut[k]]=k;
+    })
+    document.addEventListener("keydown",(e)=>{
+        if(e.target.tagName!="INPUT")
+        {
+            let key=e.key;
+            if(e.ctrlKey)
+            {
+                key="Ctrl+"+e.key;
+                if(key in key_map)
+                {
+                    e.preventDefault();
+                }
+            }
+            if(key in key_map)
+            {
+                w2ui.toolbar.click(key_map[key]);
+            }
+        }
+    })
+}

js/upload.js

@@ -0,0 +1,19 @@
+function(a,b){
+    if(!window.my_observe_upload)
+    {
+        console.log("setup upload here");
+        window.my_observe_upload = new MutationObserver(function (event) {
+            console.log(event);
+            var frame=document.querySelector("gradio-app").shadowRoot.querySelector("#sdinfframe").contentWindow.document;
+            frame.querySelector("#upload").click();
+        });
+        window.my_observe_upload_target = document.querySelector("gradio-app").shadowRoot.querySelector("#upload span");
+        window.my_observe_upload.observe(window.my_observe_upload_target, {
+            attributes: false, 
+            subtree: true,
+            childList: true, 
+            characterData: true
+        });
+    }
+    return [a,b];
+}

js/xss.js

@@ -0,0 +1,31 @@
+var setup_outpaint=function(){
+    if(!window.my_observe_outpaint)
+    {
+        console.log("setup outpaint here");
+        window.my_observe_outpaint = new MutationObserver(function (event) {
+            console.log(event);
+            let app=document.querySelector("gradio-app");
+            app=app.shadowRoot??app;
+            let frame=app.querySelector("#sdinfframe").contentWindow;
+            frame.postMessage(["outpaint", ""], "*");
+        });
+        var app=document.querySelector("gradio-app");
+        app=app.shadowRoot??app;
+        window.my_observe_outpaint_target=app.querySelector("#output span");
+        window.my_observe_outpaint.observe(window.my_observe_outpaint_target, {
+            attributes: false, 
+            subtree: true,
+            childList: true, 
+            characterData: true
+        });
+    }
+}; 
+window.config_obj={
+    resize_check: true,
+    enable_safety: true,
+    use_correction: false,
+    enable_img2img: false,
+    use_seed: false,
+    seed_val: 0,
+};
+setup_outpaint();

models/v1-inference.yaml

@@ -0,0 +1,70 @@
+model:
+  base_learning_rate: 1.0e-04
+  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "jpg"
+    cond_stage_key: "txt"
+    image_size: 64
+    channels: 4
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: crossattn
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+    use_ema: False
+
+    scheduler_config: # 10000 warmup steps
+      target: ldm.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 10000 ]
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    unet_config:
+      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: ldm.modules.encoders.modules.FrozenCLIPEmbedder

models/v1-inpainting-inference.yaml

@@ -0,0 +1,70 @@
+model:
+  base_learning_rate: 7.5e-05
+  target: ldm.models.diffusion.ddpm.LatentInpaintDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "jpg"
+    cond_stage_key: "txt"
+    image_size: 64
+    channels: 4
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: hybrid   # important
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+    finetune_keys: null
+
+    scheduler_config: # 10000 warmup steps
+      target: ldm.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 2500 ] # NOTE for resuming. use 10000 if starting from scratch
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    unet_config:
+      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 9  # 4 data + 4 downscaled image + 1 mask
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: ldm.modules.encoders.modules.FrozenCLIPEmbedder

packages.txt

@@ -0,0 +1,4 @@
+build-essential
+python3-opencv
+libopencv-dev
+cmake

perlin2d.py

@@ -0,0 +1,45 @@
+import numpy as np
+
+##########
+# https://stackoverflow.com/questions/42147776/producing-2d-perlin-noise-with-numpy/42154921#42154921
+def perlin(x, y, seed=0):
+    # permutation table
+    np.random.seed(seed)
+    p = np.arange(256, dtype=int)
+    np.random.shuffle(p)
+    p = np.stack([p, p]).flatten()
+    # coordinates of the top-left
+    xi, yi = x.astype(int), y.astype(int)
+    # internal coordinates
+    xf, yf = x - xi, y - yi
+    # fade factors
+    u, v = fade(xf), fade(yf)
+    # noise components
+    n00 = gradient(p[p[xi] + yi], xf, yf)
+    n01 = gradient(p[p[xi] + yi + 1], xf, yf - 1)
+    n11 = gradient(p[p[xi + 1] + yi + 1], xf - 1, yf - 1)
+    n10 = gradient(p[p[xi + 1] + yi], xf - 1, yf)
+    # combine noises
+    x1 = lerp(n00, n10, u)
+    x2 = lerp(n01, n11, u)  # FIX1: I was using n10 instead of n01
+    return lerp(x1, x2, v)  # FIX2: I also had to reverse x1 and x2 here
+
+
+def lerp(a, b, x):
+    "linear interpolation"
+    return a + x * (b - a)
+
+
+def fade(t):
+    "6t^5 - 15t^4 + 10t^3"
+    return 6 * t ** 5 - 15 * t ** 4 + 10 * t ** 3
+
+
+def gradient(h, x, y):
+    "grad converts h to the right gradient vector and return the dot product with (x,y)"
+    vectors = np.array([[0, 1], [0, -1], [1, 0], [-1, 0]])
+    g = vectors[h % 4]
+    return g[:, :, 0] * x + g[:, :, 1] * y
+
+
+##########

postprocess.py

@@ -0,0 +1,249 @@
+"""
+https://github.com/Trinkle23897/Fast-Poisson-Image-Editing
+MIT License
+
+Copyright (c) 2022 Jiayi Weng
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+
+import time
+import argparse
+import os
+import fpie
+from process import ALL_BACKEND, CPU_COUNT, DEFAULT_BACKEND
+from fpie.io import read_images, write_image
+from process import BaseProcessor, EquProcessor, GridProcessor
+
+from PIL import Image
+import numpy as np
+import skimage
+import skimage.measure
+import scipy
+import scipy.signal
+
+
+class PhotometricCorrection:
+    def __init__(self,quite=False):
+        self.get_parser("cli")
+        args=self.parser.parse_args(["--method","grid","-g","src","-s","a","-t","a","-o","a"])
+        args.mpi_sync_interval = getattr(args, "mpi_sync_interval", 0)
+        self.backend=args.backend
+        self.args=args
+        self.quite=quite
+        proc: BaseProcessor
+        proc = GridProcessor(
+            args.gradient,
+            args.backend,
+            args.cpu,
+            args.mpi_sync_interval,
+            args.block_size,
+            args.grid_x,
+            args.grid_y,
+        )
+        print(
+            f"[PIE]Successfully initialize PIE {args.method} solver "
+            f"with {args.backend} backend"
+        )
+        self.proc=proc
+
+    def run(self, original_image, inpainted_image, mode="mask_mode"):
+        print(f"[PIE] start")
+        if mode=="disabled":
+            return inpainted_image
+        input_arr=np.array(original_image)
+        if input_arr[:,:,-1].sum()<1:
+            return inpainted_image
+        output_arr=np.array(inpainted_image)
+        mask=input_arr[:,:,-1]
+        mask=255-mask
+        if mask.sum()<1 and mode=="mask_mode":
+            mode=""
+        if mode=="mask_mode":
+            mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
+            mask = mask.repeat(8, axis=0).repeat(8, axis=1)
+        else:
+            mask[8:-9,8:-9]=255
+        mask = mask[:,:,np.newaxis].repeat(3,axis=2)
+        nmask=mask.copy()
+        output_arr2=output_arr[:,:,0:3].copy()
+        input_arr2=input_arr[:,:,0:3].copy()
+        output_arr2[nmask<128]=0
+        input_arr2[nmask>=128]=0
+        output_arr2+=input_arr2
+        src = output_arr2[:,:,0:3]
+        tgt = src.copy()
+        proc=self.proc
+        args=self.args
+        if proc.root:
+            n = proc.reset(src, mask, tgt, (args.h0, args.w0), (args.h1, args.w1))
+        proc.sync()
+        if proc.root:
+            result = tgt
+            t = time.time()
+        if args.p == 0:
+            args.p = args.n
+
+        for i in range(0, args.n, args.p):
+            if proc.root:
+                result, err = proc.step(args.p)  # type: ignore
+                print(f"[PIE] Iter {i + args.p}, abs_err {err}")
+            else:
+                proc.step(args.p)
+
+        if proc.root:
+            dt = time.time() - t
+            print(f"[PIE] Time elapsed: {dt:.4f}s")
+            # make sure consistent with dummy process
+            return Image.fromarray(result)
+
+
+    def get_parser(self,gen_type: str) -> argparse.Namespace:
+        parser = argparse.ArgumentParser()
+        parser.add_argument(
+            "-v", "--version", action="store_true", help="show the version and exit"
+        )
+        parser.add_argument(
+            "--check-backend", action="store_true", help="print all available backends"
+        )
+        if gen_type == "gui" and "mpi" in ALL_BACKEND:
+            # gui doesn't support MPI backend
+            ALL_BACKEND.remove("mpi")
+        parser.add_argument(
+            "-b",
+            "--backend",
+            type=str,
+            choices=ALL_BACKEND,
+            default=DEFAULT_BACKEND,
+            help="backend choice",
+        )
+        parser.add_argument(
+            "-c",
+            "--cpu",
+            type=int,
+            default=CPU_COUNT,
+            help="number of CPU used",
+        )
+        parser.add_argument(
+            "-z",
+            "--block-size",
+            type=int,
+            default=1024,
+            help="cuda block size (only for equ solver)",
+        )
+        parser.add_argument(
+            "--method",
+            type=str,
+            choices=["equ", "grid"],
+            default="equ",
+            help="how to parallelize computation",
+        )
+        parser.add_argument("-s", "--source", type=str, help="source image filename")
+        if gen_type == "cli":
+            parser.add_argument(
+                "-m",
+                "--mask",
+                type=str,
+                help="mask image filename (default is to use the whole source image)",
+                default="",
+            )
+        parser.add_argument("-t", "--target", type=str, help="target image filename")
+        parser.add_argument("-o", "--output", type=str, help="output image filename")
+        if gen_type == "cli":
+            parser.add_argument(
+                "-h0", type=int, help="mask position (height) on source image", default=0
+            )
+            parser.add_argument(
+                "-w0", type=int, help="mask position (width) on source image", default=0
+            )
+            parser.add_argument(
+                "-h1", type=int, help="mask position (height) on target image", default=0
+            )
+            parser.add_argument(
+                "-w1", type=int, help="mask position (width) on target image", default=0
+            )
+        parser.add_argument(
+            "-g",
+            "--gradient",
+            type=str,
+            choices=["max", "src", "avg"],
+            default="max",
+            help="how to calculate gradient for PIE",
+        )
+        parser.add_argument(
+            "-n",
+            type=int,
+            help="how many iteration would you perfer, the more the better",
+            default=5000,
+        )
+        if gen_type == "cli":
+            parser.add_argument(
+                "-p", type=int, help="output result every P iteration", default=0
+            )
+        if "mpi" in ALL_BACKEND:
+            parser.add_argument(
+                "--mpi-sync-interval",
+                type=int,
+                help="MPI sync iteration interval",
+                default=100,
+            )
+        parser.add_argument(
+            "--grid-x", type=int, help="x axis stride for grid solver", default=8
+        )
+        parser.add_argument(
+            "--grid-y", type=int, help="y axis stride for grid solver", default=8
+        )
+        self.parser=parser
+
+if __name__ =="__main__":
+    import sys
+    import io
+    import base64
+    from PIL import Image
+    def base64_to_pil(base64_str):
+        data = base64.b64decode(str(base64_str))
+        pil = Image.open(io.BytesIO(data))
+        return pil
+
+    def pil_to_base64(out_pil):
+        out_buffer = io.BytesIO()
+        out_pil.save(out_buffer, format="PNG")
+        out_buffer.seek(0)
+        base64_bytes = base64.b64encode(out_buffer.read())
+        base64_str = base64_bytes.decode("ascii")
+        return base64_str
+    correction_func=PhotometricCorrection(quite=True)
+    while True:
+        buffer = sys.stdin.readline()
+        print(f"[PIE] suprocess {len(buffer)} {type(buffer)} ")
+        if len(buffer)==0:
+            break
+        if isinstance(buffer,str):
+            lst=buffer.strip().split(",")
+        else:
+            lst=buffer.decode("ascii").strip().split(",")
+        img0=base64_to_pil(lst[0])
+        img1=base64_to_pil(lst[1])
+        ret=correction_func.run(img0,img1,mode=lst[2])
+        ret_base64=pil_to_base64(ret)
+        if isinstance(buffer,str):
+            sys.stdout.write(f"{ret_base64}\n")
+        else:
+            sys.stdout.write(f"{ret_base64}\n".encode())
+        sys.stdout.flush()

process.py

@@ -0,0 +1,395 @@
+"""
+https://github.com/Trinkle23897/Fast-Poisson-Image-Editing
+MIT License
+
+Copyright (c) 2022 Jiayi Weng
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+import os
+from abc import ABC, abstractmethod
+from typing import Any, Optional, Tuple
+
+import numpy as np
+
+from fpie import np_solver
+
+import scipy
+import scipy.signal
+
+CPU_COUNT = os.cpu_count() or 1
+DEFAULT_BACKEND = "numpy"
+ALL_BACKEND = ["numpy"]
+
+try:
+  from fpie import numba_solver
+  ALL_BACKEND += ["numba"]
+  DEFAULT_BACKEND = "numba"
+except ImportError:
+  numba_solver = None  # type: ignore
+
+try:
+  from fpie import taichi_solver
+  ALL_BACKEND += ["taichi-cpu", "taichi-gpu"]
+  DEFAULT_BACKEND = "taichi-cpu"
+except ImportError:
+  taichi_solver = None  # type: ignore
+
+# try:
+#   from fpie import core_gcc  # type: ignore
+#   DEFAULT_BACKEND = "gcc"
+#   ALL_BACKEND.append("gcc")
+# except ImportError:
+#   core_gcc = None
+
+# try:
+#   from fpie import core_openmp  # type: ignore
+#   DEFAULT_BACKEND = "openmp"
+#   ALL_BACKEND.append("openmp")
+# except ImportError:
+#   core_openmp = None
+
+# try:
+#   from mpi4py import MPI
+
+#   from fpie import core_mpi  # type: ignore
+#   ALL_BACKEND.append("mpi")
+# except ImportError:
+#   MPI = None  # type: ignore
+#   core_mpi = None
+
+try:
+  from fpie import core_cuda  # type: ignore
+  DEFAULT_BACKEND = "cuda"
+  ALL_BACKEND.append("cuda")
+except ImportError:
+  core_cuda = None
+
+
+class BaseProcessor(ABC):
+  """API definition for processor class."""
+
+  def __init__(
+    self, gradient: str, rank: int, backend: str, core: Optional[Any]
+  ):
+    if core is None:
+      error_msg = {
+        "numpy":
+          "Please run `pip install numpy`.",
+        "numba":
+          "Please run `pip install numba`.",
+        "gcc":
+          "Please install cmake and gcc in your operating system.",
+        "openmp":
+          "Please make sure your gcc is compatible with `-fopenmp` option.",
+        "mpi":
+          "Please install MPI and run `pip install mpi4py`.",
+        "cuda":
+          "Please make sure nvcc and cuda-related libraries are available.",
+        "taichi":
+          "Please run `pip install taichi`.",
+      }
+      print(error_msg[backend.split("-")[0]])
+
+      raise AssertionError(f"Invalid backend {backend}.")
+
+    self.gradient = gradient
+    self.rank = rank
+    self.backend = backend
+    self.core = core
+    self.root = rank == 0
+
+  def mixgrad(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    if self.gradient == "src":
+      return a
+    if self.gradient == "avg":
+      return (a + b) / 2
+    # mix gradient, see Equ. 12 in PIE paper
+    mask = np.abs(a) < np.abs(b)
+    a[mask] = b[mask]
+    return a
+
+  @abstractmethod
+  def reset(
+    self,
+    src: np.ndarray,
+    mask: np.ndarray,
+    tgt: np.ndarray,
+    mask_on_src: Tuple[int, int],
+    mask_on_tgt: Tuple[int, int],
+  ) -> int:
+    pass
+
+  def sync(self) -> None:
+    self.core.sync()
+
+  @abstractmethod
+  def step(self, iteration: int) -> Optional[Tuple[np.ndarray, np.ndarray]]:
+    pass
+
+
+class EquProcessor(BaseProcessor):
+  """PIE Jacobi equation processor."""
+
+  def __init__(
+    self,
+    gradient: str = "max",
+    backend: str = DEFAULT_BACKEND,
+    n_cpu: int = CPU_COUNT,
+    min_interval: int = 100,
+    block_size: int = 1024,
+  ):
+    core: Optional[Any] = None
+    rank = 0
+
+    if backend == "numpy":
+      core = np_solver.EquSolver()
+    elif backend == "numba" and numba_solver is not None:
+      core = numba_solver.EquSolver()
+    elif backend == "gcc":
+      core = core_gcc.EquSolver()
+    elif backend == "openmp" and core_openmp is not None:
+      core = core_openmp.EquSolver(n_cpu)
+    elif backend == "mpi" and core_mpi is not None:
+      core = core_mpi.EquSolver(min_interval)
+      rank = MPI.COMM_WORLD.Get_rank()
+    elif backend == "cuda" and core_cuda is not None:
+      core = core_cuda.EquSolver(block_size)
+    elif backend.startswith("taichi") and taichi_solver is not None:
+      core = taichi_solver.EquSolver(backend, n_cpu, block_size)
+
+    super().__init__(gradient, rank, backend, core)
+
+  def mask2index(
+    self, mask: np.ndarray
+  ) -> Tuple[np.ndarray, int, np.ndarray, np.ndarray]:
+    x, y = np.nonzero(mask)
+    max_id = x.shape[0] + 1
+    index = np.zeros((max_id, 3))
+    ids = self.core.partition(mask)
+    ids[mask == 0] = 0  # reserve id=0 for constant
+    index = ids[x, y].argsort()
+    return ids, max_id, x[index], y[index]
+
+  def reset(
+    self,
+    src: np.ndarray,
+    mask: np.ndarray,
+    tgt: np.ndarray,
+    mask_on_src: Tuple[int, int],
+    mask_on_tgt: Tuple[int, int],
+  ) -> int:
+    assert self.root
+    # check validity
+    # assert 0 <= mask_on_src[0] and 0 <= mask_on_src[1]
+    # assert mask_on_src[0] + mask.shape[0] <= src.shape[0]
+    # assert mask_on_src[1] + mask.shape[1] <= src.shape[1]
+    # assert mask_on_tgt[0] + mask.shape[0] <= tgt.shape[0]
+    # assert mask_on_tgt[1] + mask.shape[1] <= tgt.shape[1]
+
+    if len(mask.shape) == 3:
+      mask = mask.mean(-1)
+    mask = (mask >= 128).astype(np.int32)
+
+    # zero-out edge
+    mask[0] = 0
+    mask[-1] = 0
+    mask[:, 0] = 0
+    mask[:, -1] = 0
+
+    x, y = np.nonzero(mask)
+    x0, x1 = x.min() - 1, x.max() + 2
+    y0, y1 = y.min() - 1, y.max() + 2
+    mask_on_src = (x0 + mask_on_src[0], y0 + mask_on_src[1])
+    mask_on_tgt = (x0 + mask_on_tgt[0], y0 + mask_on_tgt[1])
+    mask = mask[x0:x1, y0:y1]
+    ids, max_id, index_x, index_y = self.mask2index(mask)
+
+    src_x, src_y = index_x + mask_on_src[0], index_y + mask_on_src[1]
+    tgt_x, tgt_y = index_x + mask_on_tgt[0], index_y + mask_on_tgt[1]
+
+    src_C = src[src_x, src_y].astype(np.float32)
+    src_U = src[src_x - 1, src_y].astype(np.float32)
+    src_D = src[src_x + 1, src_y].astype(np.float32)
+    src_L = src[src_x, src_y - 1].astype(np.float32)
+    src_R = src[src_x, src_y + 1].astype(np.float32)
+    tgt_C = tgt[tgt_x, tgt_y].astype(np.float32)
+    tgt_U = tgt[tgt_x - 1, tgt_y].astype(np.float32)
+    tgt_D = tgt[tgt_x + 1, tgt_y].astype(np.float32)
+    tgt_L = tgt[tgt_x, tgt_y - 1].astype(np.float32)
+    tgt_R = tgt[tgt_x, tgt_y + 1].astype(np.float32)
+
+    grad = self.mixgrad(src_C - src_L, tgt_C - tgt_L) \
+      + self.mixgrad(src_C - src_R, tgt_C - tgt_R) \
+      + self.mixgrad(src_C - src_U, tgt_C - tgt_U) \
+      + self.mixgrad(src_C - src_D, tgt_C - tgt_D)
+
+    A = np.zeros((max_id, 4), np.int32)
+    X = np.zeros((max_id, 3), np.float32)
+    B = np.zeros((max_id, 3), np.float32)
+
+    X[1:] = tgt[index_x + mask_on_tgt[0], index_y + mask_on_tgt[1]]
+    # four-way
+    A[1:, 0] = ids[index_x - 1, index_y]
+    A[1:, 1] = ids[index_x + 1, index_y]
+    A[1:, 2] = ids[index_x, index_y - 1]
+    A[1:, 3] = ids[index_x, index_y + 1]
+    B[1:] = grad
+    m = (mask[index_x - 1, index_y] == 0).astype(float).reshape(-1, 1)
+    B[1:] += m * tgt[index_x + mask_on_tgt[0] - 1, index_y + mask_on_tgt[1]]
+    m = (mask[index_x, index_y - 1] == 0).astype(float).reshape(-1, 1)
+    B[1:] += m * tgt[index_x + mask_on_tgt[0], index_y + mask_on_tgt[1] - 1]
+    m = (mask[index_x, index_y + 1] == 0).astype(float).reshape(-1, 1)
+    B[1:] += m * tgt[index_x + mask_on_tgt[0], index_y + mask_on_tgt[1] + 1]
+    m = (mask[index_x + 1, index_y] == 0).astype(float).reshape(-1, 1)
+    B[1:] += m * tgt[index_x + mask_on_tgt[0] + 1, index_y + mask_on_tgt[1]]
+
+    self.tgt = tgt.copy()
+    self.tgt_index = (index_x + mask_on_tgt[0], index_y + mask_on_tgt[1])
+    self.core.reset(max_id, A, X, B)
+    return max_id
+
+  def step(self, iteration: int) -> Optional[Tuple[np.ndarray, np.ndarray]]:
+    result = self.core.step(iteration)
+    if self.root:
+      x, err = result
+      self.tgt[self.tgt_index] = x[1:]
+      return self.tgt, err
+    return None
+
+
+class GridProcessor(BaseProcessor):
+  """PIE grid processor."""
+
+  def __init__(
+    self,
+    gradient: str = "max",
+    backend: str = DEFAULT_BACKEND,
+    n_cpu: int = CPU_COUNT,
+    min_interval: int = 100,
+    block_size: int = 1024,
+    grid_x: int = 8,
+    grid_y: int = 8,
+  ):
+    core: Optional[Any] = None
+    rank = 0
+
+    if backend == "numpy":
+      core = np_solver.GridSolver()
+    elif backend == "numba" and numba_solver is not None:
+      core = numba_solver.GridSolver()
+    elif backend == "gcc":
+      core = core_gcc.GridSolver(grid_x, grid_y)
+    elif backend == "openmp" and core_openmp is not None:
+      core = core_openmp.GridSolver(grid_x, grid_y, n_cpu)
+    elif backend == "mpi" and core_mpi is not None:
+      core = core_mpi.GridSolver(min_interval)
+      rank = MPI.COMM_WORLD.Get_rank()
+    elif backend == "cuda" and core_cuda is not None:
+      core = core_cuda.GridSolver(grid_x, grid_y)
+    elif backend.startswith("taichi") and taichi_solver is not None:
+      core = taichi_solver.GridSolver(
+        grid_x, grid_y, backend, n_cpu, block_size
+      )
+
+    super().__init__(gradient, rank, backend, core)
+
+  def reset(
+    self,
+    src: np.ndarray,
+    mask: np.ndarray,
+    tgt: np.ndarray,
+    mask_on_src: Tuple[int, int],
+    mask_on_tgt: Tuple[int, int],
+  ) -> int:
+    assert self.root
+    # check validity
+    # assert 0 <= mask_on_src[0] and 0 <= mask_on_src[1]
+    # assert mask_on_src[0] + mask.shape[0] <= src.shape[0]
+    # assert mask_on_src[1] + mask.shape[1] <= src.shape[1]
+    # assert mask_on_tgt[0] + mask.shape[0] <= tgt.shape[0]
+    # assert mask_on_tgt[1] + mask.shape[1] <= tgt.shape[1]
+
+    if len(mask.shape) == 3:
+      mask = mask.mean(-1)
+    mask = (mask >= 128).astype(np.int32)
+
+    # zero-out edge
+    mask[0] = 0
+    mask[-1] = 0
+    mask[:, 0] = 0
+    mask[:, -1] = 0
+
+    x, y = np.nonzero(mask)
+    x0, x1 = x.min() - 1, x.max() + 2
+    y0, y1 = y.min() - 1, y.max() + 2
+    mask = mask[x0:x1, y0:y1]
+    max_id = np.prod(mask.shape)
+
+    src_crop = src[mask_on_src[0] + x0:mask_on_src[0] + x1,
+                   mask_on_src[1] + y0:mask_on_src[1] + y1].astype(np.float32)
+    tgt_crop = tgt[mask_on_tgt[0] + x0:mask_on_tgt[0] + x1,
+                   mask_on_tgt[1] + y0:mask_on_tgt[1] + y1].astype(np.float32)
+    grad = np.zeros([*mask.shape, 3], np.float32)
+    grad[1:] += self.mixgrad(
+      src_crop[1:] - src_crop[:-1], tgt_crop[1:] - tgt_crop[:-1]
+    )
+    grad[:-1] += self.mixgrad(
+      src_crop[:-1] - src_crop[1:], tgt_crop[:-1] - tgt_crop[1:]
+    )
+    grad[:, 1:] += self.mixgrad(
+      src_crop[:, 1:] - src_crop[:, :-1], tgt_crop[:, 1:] - tgt_crop[:, :-1]
+    )
+    grad[:, :-1] += self.mixgrad(
+      src_crop[:, :-1] - src_crop[:, 1:], tgt_crop[:, :-1] - tgt_crop[:, 1:]
+    )
+
+    grad[mask == 0] = 0
+    if True:
+        kernel = [[1] * 3 for _ in range(3)]
+        nmask = mask.copy()
+        nmask[nmask > 0] = 1
+        res = scipy.signal.convolve2d(
+            nmask, kernel, mode="same", boundary="fill", fillvalue=1
+        )
+        res[nmask < 1] = 0
+        res[res == 9] = 0
+        res[res > 0] = 1
+        grad[res>0]=0
+        # ylst, xlst = res.nonzero()
+        # for y, x in zip(ylst, xlst):
+        #     grad[y,x]=0
+            # for yi in range(-1,2):
+                # for xi in range(-1,2):
+                    # grad[y+yi,x+xi]=0
+    self.x0 = mask_on_tgt[0] + x0
+    self.x1 = mask_on_tgt[0] + x1
+    self.y0 = mask_on_tgt[1] + y0
+    self.y1 = mask_on_tgt[1] + y1
+    self.tgt = tgt.copy()
+    self.core.reset(max_id, mask, tgt_crop, grad)
+    return max_id
+
+  def step(self, iteration: int) -> Optional[Tuple[np.ndarray, np.ndarray]]:
+    result = self.core.step(iteration)
+    if self.root:
+      tgt, err = result
+      self.tgt[self.x0:self.x1, self.y0:self.y1] = tgt
+      return self.tgt, err
+    return None

requirements.txt

@@ -0,0 +1,17 @@
+--extra-index-url https://download.pytorch.org/whl/nightly/cu117
+imageio==2.19.5
+imageio-ffmpeg==0.4.7
+numpy==1.22.4
+opencv-python-headless==4.6.0.66
+torch[dynamo]>=0.0.dev0
+torchvision
+Pillow
+scipy
+scikit-image
+diffusers==0.9.0
+transformers
+ftfy
+fpie
+accelerate
+ninja
+setuptools==59.8.0

utils.py

@@ -0,0 +1,217 @@
+from PIL import Image
+from PIL import ImageFilter
+import cv2
+import numpy as np
+import scipy
+import scipy.signal
+from scipy.spatial import cKDTree
+
+import os
+from perlin2d import *
+
+patch_match_compiled = True
+
+try:
+    from PyPatchMatch import patch_match
+except Exception as e:
+    try:
+        import patch_match
+    except Exception as e:
+        patch_match_compiled = False
+
+try:
+    patch_match
+except NameError:
+    print("patch_match compiling failed, will fall back to edge_pad")
+    patch_match_compiled = False
+
+
+
+
+def edge_pad(img, mask, mode=1):
+    if mode == 0:
+        nmask = mask.copy()
+        nmask[nmask > 0] = 1
+        res0 = 1 - nmask
+        res1 = nmask
+        p0 = np.stack(res0.nonzero(), axis=0).transpose()
+        p1 = np.stack(res1.nonzero(), axis=0).transpose()
+        min_dists, min_dist_idx = cKDTree(p1).query(p0, 1)
+        loc = p1[min_dist_idx]
+        for (a, b), (c, d) in zip(p0, loc):
+            img[a, b] = img[c, d]
+    elif mode == 1:
+        record = {}
+        kernel = [[1] * 3 for _ in range(3)]
+        nmask = mask.copy()
+        nmask[nmask > 0] = 1
+        res = scipy.signal.convolve2d(
+            nmask, kernel, mode="same", boundary="fill", fillvalue=1
+        )
+        res[nmask < 1] = 0
+        res[res == 9] = 0
+        res[res > 0] = 1
+        ylst, xlst = res.nonzero()
+        queue = [(y, x) for y, x in zip(ylst, xlst)]
+        # bfs here
+        cnt = res.astype(np.float32)
+        acc = img.astype(np.float32)
+        step = 1
+        h = acc.shape[0]
+        w = acc.shape[1]
+        offset = [(1, 0), (-1, 0), (0, 1), (0, -1)]
+        while queue:
+            target = []
+            for y, x in queue:
+                val = acc[y][x]
+                for yo, xo in offset:
+                    yn = y + yo
+                    xn = x + xo
+                    if 0 <= yn < h and 0 <= xn < w and nmask[yn][xn] < 1:
+                        if record.get((yn, xn), step) == step:
+                            acc[yn][xn] = acc[yn][xn] * cnt[yn][xn] + val
+                            cnt[yn][xn] += 1
+                            acc[yn][xn] /= cnt[yn][xn]
+                            if (yn, xn) not in record:
+                                record[(yn, xn)] = step
+                                target.append((yn, xn))
+            step += 1
+            queue = target
+        img = acc.astype(np.uint8)
+    else:
+        nmask = mask.copy()
+        ylst, xlst = nmask.nonzero()
+        yt, xt = ylst.min(), xlst.min()
+        yb, xb = ylst.max(), xlst.max()
+        content = img[yt : yb + 1, xt : xb + 1]
+        img = np.pad(
+            content,
+            ((yt, mask.shape[0] - yb - 1), (xt, mask.shape[1] - xb - 1), (0, 0)),
+            mode="edge",
+        )
+    return img, mask
+
+
+def perlin_noise(img, mask):
+    lin = np.linspace(0, 5, mask.shape[0], endpoint=False)
+    x, y = np.meshgrid(lin, lin)
+    avg = img.mean(axis=0).mean(axis=0)
+    # noise=[((perlin(x, y)+1)*128+avg[i]).astype(np.uint8) for i in range(3)]
+    noise = [((perlin(x, y) + 1) * 0.5 * 255).astype(np.uint8) for i in range(3)]
+    noise = np.stack(noise, axis=-1)
+    # mask=skimage.measure.block_reduce(mask,(8,8),np.min)
+    # mask=mask.repeat(8, axis=0).repeat(8, axis=1)
+    # mask_image=Image.fromarray(mask)
+    # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 4))
+    # mask=np.array(mask_image)
+    nmask = mask.copy()
+    # nmask=nmask/255.0
+    nmask[mask > 0] = 1
+    img = nmask[:, :, np.newaxis] * img + (1 - nmask[:, :, np.newaxis]) * noise
+    # img=img.astype(np.uint8)
+    return img, mask
+
+
+def gaussian_noise(img, mask):
+    noise = np.random.randn(mask.shape[0], mask.shape[1], 3)
+    noise = (noise + 1) / 2 * 255
+    noise = noise.astype(np.uint8)
+    nmask = mask.copy()
+    nmask[mask > 0] = 1
+    img = nmask[:, :, np.newaxis] * img + (1 - nmask[:, :, np.newaxis]) * noise
+    return img, mask
+
+
+def cv2_telea(img, mask):
+    ret = cv2.inpaint(img, 255 - mask, 5, cv2.INPAINT_TELEA)
+    return ret, mask
+
+
+def cv2_ns(img, mask):
+    ret = cv2.inpaint(img, 255 - mask, 5, cv2.INPAINT_NS)
+    return ret, mask
+
+
+def patch_match_func(img, mask):
+    ret = patch_match.inpaint(img, mask=255 - mask, patch_size=3)
+    return ret, mask
+
+
+def mean_fill(img, mask):
+    avg = img.mean(axis=0).mean(axis=0)
+    img[mask < 1] = avg
+    return img, mask
+
+def g_diffuser(img,mask):
+    return img, mask
+
+def dummy_fill(img,mask):
+    return img,mask
+functbl = {
+    "gaussian": gaussian_noise,
+    "perlin": perlin_noise,
+    "edge_pad": edge_pad,
+    "patchmatch": patch_match_func if patch_match_compiled else edge_pad,
+    "cv2_ns": cv2_ns,
+    "cv2_telea": cv2_telea,
+    "g_diffuser": g_diffuser,
+    "g_diffuser_lib": dummy_fill,
+}
+
+try:
+    from postprocess import PhotometricCorrection
+    correction_func = PhotometricCorrection()
+except Exception as e:
+    print(e, "so PhotometricCorrection is disabled")
+    class DummyCorrection:
+        def __init__(self):
+            self.backend=""
+            pass
+        def run(self,a,b,**kwargs):
+            return b
+    correction_func=DummyCorrection()
+
+if "taichi" in correction_func.backend:
+    import sys
+    import io
+    import base64
+    from PIL import Image
+    def base64_to_pil(base64_str):
+        data = base64.b64decode(str(base64_str))
+        pil = Image.open(io.BytesIO(data))
+        return pil
+
+    def pil_to_base64(out_pil):
+        out_buffer = io.BytesIO()
+        out_pil.save(out_buffer, format="PNG")
+        out_buffer.seek(0)
+        base64_bytes = base64.b64encode(out_buffer.read())
+        base64_str = base64_bytes.decode("ascii")
+        return base64_str
+    from subprocess import Popen, PIPE, STDOUT
+    class SubprocessCorrection:
+        def __init__(self):
+            self.backend=correction_func.backend
+            self.child= Popen(["python", "postprocess.py"], stdin=PIPE, stdout=PIPE, stderr=STDOUT)
+        def run(self,img_input,img_inpainted,mode):
+            if mode=="disabled":
+                return img_inpainted
+            base64_str_input = pil_to_base64(img_input)
+            base64_str_inpainted = pil_to_base64(img_inpainted)
+            try:
+                if self.child.poll():
+                    self.child= Popen(["python", "postprocess.py"], stdin=PIPE, stdout=PIPE, stderr=STDOUT)
+                self.child.stdin.write(f"{base64_str_input},{base64_str_inpainted},{mode}\n".encode())
+                self.child.stdin.flush()
+                out = self.child.stdout.readline()
+                base64_str=out.decode().strip()
+                while base64_str and base64_str[0]=="[":
+                    print(base64_str)
+                    out = self.child.stdout.readline()
+                    base64_str=out.decode().strip()
+                ret=base64_to_pil(base64_str)
+            except:
+                print("[PIE] not working, photometric correction is disabled")
+                ret=img_inpainted
+            return ret
+    correction_func = SubprocessCorrection()