pushin to huggin

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: [email protected]
+
+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 <[email protected]>
+#
+# 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  : [email protected]
+# 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  : [email protected]
+# 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  : [email protected]
+# 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 <[email protected]>
+#
+# Distributed under terms of the MIT license.
+#
+
+make clean && make

app.py

@@ -0,0 +1,1262 @@
+import io
+import base64
+import os
+import sys
+
+import numpy as np
+import torch
+from torch import autocast
+import diffusers
+import requests
+
+
+assert tuple(map(int,diffusers.__version__.split(".")))  >= (0,9,0), "Please upgrade diffusers to 0.9.0"
+
+from diffusers.configuration_utils import FrozenDict
+from diffusers import (
+    StableDiffusionPipeline,
+    StableDiffusionInpaintPipeline,
+    StableDiffusionImg2ImgPipeline,
+    StableDiffusionInpaintPipelineLegacy,
+    DDIMScheduler,
+    LMSDiscreteScheduler,
+    DiffusionPipeline,
+    StableDiffusionUpscalePipeline,
+    DPMSolverMultistepScheduler,
+    PNDMScheduler,
+)
+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
+from utils import *
+
+# load environment variables from the .env file
+if os.path.exists(".env"):
+    with open(".env") as f:
+        for line in f:
+            if line.startswith("#") or not line.strip():
+                continue
+            name, value = line.strip().split("=", 1)
+            os.environ[name] = value
+
+
+access_token = os.environ.get("HF_ACCESS_TOKEN")
+print("access_token from HF 1:", access_token)
+
+
+
+def query(payload, model_id, api_token):
+	headers = {"Authorization": f"Bearer {api_token}"}
+	API_URL = f"https://api-inference.huggingface.co/models/{model_id}"
+	response = requests.post(API_URL, headers=headers, json=payload)
+	return response.json()
+
+model_id = "stabilityai/stable-diffusion-2-inpainting"
+api_token = "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" # get yours at hf.co/settings/tokens
+data = query("The goal of life is [MASK].", model_id, api_token)
+
+
+
+
+# def get_latest_image_url(database_url):
+#     response = requests.get(f"{database_url}/latestImage.json")
+#     latest_image_data = response.json()
+#     image_url = latest_image_data['downloadURL']
+#     image_name = latest_image_data['fileName']
+#     return image_url, image_name
+
+# database_url = 'https://nyucapstone-7c22c-default-rtdb.firebaseio.com'
+# latest_image_url, latest_image_name = get_latest_image_url(database_url)
+# print(f"Latest image URL: {latest_image_url}")
+# print(f"Latest image name: {latest_image_name}")
+
+try:
+    abspath = os.path.abspath(__file__)
+    dirname = os.path.dirname(abspath)
+    os.chdir(dirname)
+except:
+    pass
+
+try:
+    from interrogate import Interrogator
+except:
+    Interrogator = DummyInterrogator
+
+USE_NEW_DIFFUSERS = True
+RUN_IN_SPACE = "RUN_IN_HG_SPACE" in os.environ
+
+
+class ModelChoice(Enum):
+    INPAINTING = "stablediffusion-inpainting"
+    INPAINTING2 = "stablediffusion-2-inpainting"
+    INPAINTING_IMG2IMG = "stablediffusion-inpainting+img2img-1.5"
+    MODEL_2_1 = "stablediffusion-2.1"
+    MODEL_2_0_V = "stablediffusion-2.0v"
+    MODEL_2_0 = "stablediffusion-2.0"
+    MODEL_1_5 = "stablediffusion-1.5"
+    MODEL_1_4 = "stablediffusion-1.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:
+    import onnxruntime
+    onnx_available = True
+    onnx_providers = ["CUDAExecutionProvider", "DmlExecutionProvider", "OpenVINOExecutionProvider", 'CPUExecutionProvider']
+    available_providers = onnxruntime.get_available_providers()
+    onnx_providers = [item for item in onnx_providers if item in available_providers]
+except:
+    onnx_available = False
+    onnx_providers = []
+
+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"
+
+if device != "cuda":
+    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: 780px" 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("--lowvram", action="store_true", help="using lowvram mode")
+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__":
+    args = parser.parse_args()
+else:
+    args = parser.parse_args(["--debug"])
+# args = parser.parse_args(["--debug"])
+if args.auth is not None:
+    args.auth = tuple(args.auth)
+
+model = {}
+
+
+def get_token():
+    token = "{access_token}"
+    if os.path.exists(".token"):
+        with open(".token", "r") as f:
+            token = f.read()
+            print("get_token called", token)
+    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, "PNDM": None}
+
+
+class StableDiffusionInpaint:
+    def __init__(
+        self, token: str = "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", model_name: str = "", model_path: str = "", **kwargs,
+    ):
+        self.token = token
+        original_checkpoint = False
+        # if device == "cpu" and onnx_available:
+        #     from diffusers import OnnxStableDiffusionInpaintPipeline
+        #     inpaint = OnnxStableDiffusionInpaintPipeline.from_pretrained(
+        #         model_name,
+        #         revision="onnx",
+        #         provider=onnx_providers[0] if onnx_providers else None
+        #         )
+        # else:
+        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")
+        if device == "cuda" and not args.fp32:
+            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" and not args.fp32:
+                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" and not args.fp32:
+                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=access_token, vae=vae
+                )
+                print(f"access_token from HF:", access_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)
+        # 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["PNDM"] = prepare_scheduler(
+            PNDMScheduler(
+                beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
+                skip_prk_steps=True
+            )
+        )
+        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 or args.lowvram:
+                inpaint.enable_attention_slicing()
+                inpaint.enable_sequential_cpu_offload()
+        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 or self.safety_checker is None:
+                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:
+            if fill_mode == "g_diffuser":
+                mask = 255 - mask
+                mask = mask[:, :, np.newaxis].repeat(3, axis=2)
+                img, mask = functbl[fill_mode](img, mask)
+            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:
+                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
+#         if device=="cpu" and onnx_available:
+#             from diffusers import OnnxStableDiffusionPipeline, OnnxStableDiffusionInpaintPipelineLegacy, OnnxStableDiffusionImg2ImgPipeline
+#             text2img = OnnxStableDiffusionPipeline.from_pretrained(
+#                 model_name,
+#                 revision="onnx",
+#                 provider=onnx_providers[0] if onnx_providers else None
+#                 )
+#             inpaint = OnnxStableDiffusionInpaintPipelineLegacy(
+#                     vae_encoder=text2img.vae_encoder,
+#                     vae_decoder=text2img.vae_decoder,
+#                     text_encoder=text2img.text_encoder,
+#                     tokenizer=text2img.tokenizer,
+#                     unet=text2img.unet,
+#                     scheduler=text2img.scheduler,
+#                     safety_checker=text2img.safety_checker,
+#                     feature_extractor=text2img.feature_extractor,
+#                 )
+#             img2img = OnnxStableDiffusionImg2ImgPipeline(
+#                 vae_encoder=text2img.vae_encoder,
+#                 vae_decoder=text2img.vae_decoder,
+#                 text_encoder=text2img.text_encoder,
+#                 tokenizer=text2img.tokenizer,
+#                 unet=text2img.unet,
+#                 scheduler=text2img.scheduler,
+#                 safety_checker=text2img.safety_checker,
+#                 feature_extractor=text2img.feature_extractor,
+#             )
+#         else:
+#             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")
+#             if device == "cuda" and not args.fp32:
+#                 vae.to(torch.float16)
+#             if original_checkpoint:
+#                 print(f"Converting & Loading {model_path}")
+#                 from convert_checkpoint import convert_checkpoint
+
+#                 pipe = convert_checkpoint(model_path)
+#                 if device == "cuda" and not args.fp32:
+#                     pipe.to(torch.float16)
+#                 text2img = StableDiffusionPipeline(
+#                     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" and not args.fp32:
+#                     text2img = StableDiffusionPipeline.from_pretrained(
+#                         model_name,
+#                         revision="fp16",
+#                         torch_dtype=torch.float16,
+#                         use_auth_token=token,
+#                         vae=vae,
+#                     )
+#                 else:
+#                     text2img = StableDiffusionPipeline.from_pretrained(
+#                         model_name, use_auth_token=token, vae=vae
+#                     )
+#             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" and not args.fp32:
+#                     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,
+#                         vae=vae,
+#                     ).to(device)
+#                 text2img_unet.to(device)
+#                 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)
+#             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)
+#         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["PNDM"] = prepare_scheduler(
+#             PNDMScheduler(
+#                 beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
+#                 skip_prk_steps=True
+#             )
+#         )
+#         scheduler_dict["DPM"] = prepare_scheduler(
+#             DPMSolverMultistepScheduler.from_config(text2img.scheduler.config)
+#         )
+#         self.safety_checker = text2img.safety_checker
+#         save_token(token)
+#         try:
+#             total_memory = torch.cuda.get_device_properties(0).total_memory // (
+#                 1024 ** 3
+#             )
+#             if total_memory <= 5 or args.lowvram:
+#                 inpaint.enable_attention_slicing()
+#                 inpaint.enable_sequential_cpu_offload()
+#                 if inpainting_model:
+#                     text2img.enable_attention_slicing()
+#                     text2img.enable_sequential_cpu_offload()
+#         except:
+#             pass
+#         self.text2img = text2img
+#         self.inpaint = inpaint
+#         self.img2img = img2img
+#         if True:
+#             self.unified = inpaint
+#         else:
+#             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 or self.safety_checker is None:
+#                 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,
+#                     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 = functbl[fill_mode](img, mask)
+#                 extra_kwargs["strength"] = 1.0
+#                 extra_kwargs["out_mask"] = Image.fromarray(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)
+#                 inpaint_func = inpaint
+#             init_image = Image.fromarray(img)
+#             mask_image = Image.fromarray(mask)
+#             # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
+#             input_image = init_image.resize(
+#                 (process_width, process_height), resample=SAMPLING_MODE
+#             )
+#             if self.inpainting_model:
+#                 images = inpaint_func(
+#                     prompt=prompt,
+#                     image=input_image,
+#                     width=process_width,
+#                     height=process_height,
+#                     mask_image=mask_image.resize((process_width, process_height)),
+#                     **extra_kwargs,
+#                 )["images"]
+#             else:
+#                 extra_kwargs["strength"] = strength
+#                 if True:
+#                     images = inpaint_func(
+#                         prompt=prompt,
+#                         image=input_image,
+#                         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="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", model_choice="", model_path=""):
+    if "model" not in model:
+        model_name = ""
+        if args.local_model:
+            print(f"Using local_model: {args.local_model}")
+            model_path = args.local_model
+        elif args.remote_model:
+            print(f"Using remote_model: {args.remote_model}")
+            model_name = args.remote_model
+        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.INPAINTING2.value:
+            if len(model_name) < 1:
+                model_name = "stabilityai/stable-diffusion-2-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, 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"
+                )
+                if model_choice == ModelChoice.MODEL_2_0.value:
+                    model_name = "stabilityai/stable-diffusion-2-base"
+                elif model_choice == ModelChoice.MODEL_2_0_V.value:
+                    model_name = "stabilityai/stable-diffusion-2"
+                elif model_choice == ModelChoice.MODEL_2_1.value:
+                    model_name = "stabilityai/stable-diffusion-2-1-base"
+            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,
+    interrogate_mode,
+    state,
+):
+    data = base64.b64decode(str(sel_buffer_str))
+    pil = Image.open(io.BytesIO(data))
+    if interrogate_mode:
+        if "interrogator" not in model:
+            model["interrogator"] = Interrogator()
+        interrogator = model["interrogator"]
+        # possible point to integrate
+        img = np.array(pil)[:, :, 0:3]
+        mask = np.array(pil)[:, :, -1]
+        x, y = np.nonzero(mask)
+        if len(x) > 0:
+            x0, x1 = x.min(), x.max() + 1
+            y0, y1 = y.min(), y.max() + 1
+            img = img[x0:x1, y0:y1, :]
+        pil = Image.fromarray(img)
+        interrogate_ret = interrogator.interrogate(pil)
+        return (
+            gr.update(value=",".join([sel_buffer_str]),),
+            gr.update(label="Prompt", value=interrogate_ret),
+            state,
+        )
+    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_IMG2IMG.value,
+    )
+
+blocks = gr.Blocks(
+    title="StableDiffusion-Infinity",
+    css="""
+.tabs {
+margin-top: 0rem;
+margin-bottom: 0rem;
+}
+#markdown {
+min-height: 0rem;
+}
+""",
+theme=gr.themes.Soft()
+)
+model_path_input_val = ""
+with blocks as demo:
+    # title
+    title = gr.Markdown(
+        """
+    stanley capstone
+    """,
+        elem_id="markdown",
+    )
+    # github logo
+    github_logo = gr.HTML(
+        """
+        <a href="https://github.com/stanleywalker1/capstone-studio-2">
+            <svg xmlns="http://www.w3.org/2000/svg" width="24" height="24" viewBox="0 0 24 24"><path d="M12 0c-6.626 0-12 5.373-12 12 0 5.302 3.438 9.8 8.207 11.387.599.111.793-.261.793-.577v-2.234c-3.338.726-4.033-1.416-4.033-1.416-.546-1.387-1.333-1.756-1.333-1.756-1.089-.745.083-.729.083-.729 1.205.084 1.839 1.237 1.839 1.237 1.07 1.834 2.807 1.304 3.492.997.107-.775.418-1.305.762-1.604-2.665-.305-5.467-1.334-5.467-5.931 0-1.311.469-2.381 1.236-3.221-.124-.303-.535-1.524.117-3.176 0 0 1.008-.322 3.301 1.23.957-.266 1.983-.399 3.003-.404 1.02.005 2.047.138 3.006.404 2.291-1.552 3.297-1.23 3.297-1.23.653 1.653.242 2.874.118 3.176.77.84 1.235 1.911 1.235 3.221 0 4.609-2.807 5.624-5.479 5.921.43.372.823 1.102.823 2.222v3.293c0 .319.192.694.801.576 4.765-1.589 8.199-6.086 8.199-11.386 0-6.627-5.373-12-12-12z" fill="white"/></svg>
+        </a>
+        """
+    )
+    # frame
+    frame = gr.HTML(test(2), visible=RUN_IN_SPACE)
+    # setup
+
+    setup_button = gr.Button("Click to Start", variant="primary")
+   
+
+    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")
+
+        sd_prompt = gr.Textbox(
+            label="Prompt", placeholder="input your prompt here!", lines=2
+        )
+        with gr.Accordion("developer tools", open=True):
+            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 type here",
+                        choices=model_choices_lst,
+                        value=ModelChoice.INPAINTING.value if onnx_available else ModelChoice.INPAINTING2.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=700,
+                        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",
+                    )
+                with gr.Column(scale=3, min_width=270):
+                    init_mode = gr.Dropdown(
+                        label="Init Mode",
+                        choices=[
+                            "patchmatch",
+                            "edge_pad",
+                            "cv2_ns",
+                            "cv2_telea",
+                            "perlin",
+                            "gaussian",
+                            "g_diffuser",
+                        ],
+                        value="patchmatch",
+                        type="value",
+                    )
+                    postprocess_check = gr.Radio(
+                        label="Photometric Correction Mode",
+                        choices=["disabled", "mask_mode", "border_mode",],
+                        value="disabled",
+                        type="value",
+                    )
+                    # canvas control
+
+                with gr.Column(scale=3, min_width=270):
+                    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=1.0,
+                                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)
+                
+            model_path_input = gr.Textbox(
+                value=model_path_input_val,
+                label="Custom Model Path (You have to select a correct model type for your local model)",
+                placeholder="Ignore this if you are not using Docker",
+                elem_id="model_path_input",
+            )
+    
+    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)
+    interrogate_check = gr.Checkbox(label="Interrogate", value=False, 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))}
+            if model_choice in [
+                ModelChoice.INPAINTING.value,
+                ModelChoice.INPAINTING_IMG2IMG.value,
+                ModelChoice.INPAINTING2.value,
+            ]:
+                init_val = "cv2_ns"
+            else:
+                init_val = "patchmatch"
+            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),
+                init_mode: gr.update(value=init_val),
+            }
+
+        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,
+                init_mode,
+            ],
+            _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,
+            interrogate_check,
+            model_output_state,
+        ],
+        outputs=[model_output, sd_prompt, model_output_state],
+        _js=proceed_button_js,
+    )
+    # cancel button can also remove error overlay
+    if tuple(map(int,gr.__version__.split("."))) >= (3,6):
+        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.pop("lowvram", None)
+launch_kwargs.update(launch_extra_kwargs)
+try:
+    import google.colab
+
+    launch_kwargs["debug"] = True
+except:
+    pass
+
+if RUN_IN_SPACE:
+    demo.launch(share=True)
+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,718 @@
+import base64
+import json
+import io
+import numpy as np
+from PIL import Image
+from pyodide import to_js, create_proxy
+from pyodide.http import pyfetch
+
+import gc
+from js import (
+    console,
+    document,
+    devicePixelRatio,
+    ImageData,
+    Uint8ClampedArray,
+    CanvasRenderingContext2D as Context2d,
+    requestAnimationFrame,
+    update_overlay,
+    setup_overlay,
+    window,
+    alert,
+    fetch,
+    console
+)
+
+PAINT_SELECTION = "selection"
+IMAGE_SELECTION = "canvas"
+BRUSH_SELECTION = "eraser"
+NOP_MODE = 0
+PAINT_MODE = 1
+IMAGE_MODE = 2
+BRUSH_MODE = 3
+
+
+from js import Image as JsImage
+
+
+# async def fetch_latest_image_url(database_url):
+#         console.log("fetch_latest_image called from canvas")
+#         # different methods to call
+#         response = await fetch(f"{database_url}/latestImage.json")
+#         console.log(f"response status: {response.status}, status text: {response.statusText}")
+        
+#         latest_image_data = await response.json()
+#         latest_image_data = latest_image_data.to_py()
+
+#         image_url = latest_image_data["downloadURL"]
+#         image_name = latest_image_data["fileName"]
+#         console.log(f"Latest image URL from canvas: {image_url}")
+#         console.log(f"Latest image name from canvas: {image_name}")
+
+#         # Fetch the image data as ArrayBuffer
+#         image_response = await fetch(image_url)
+#         image_data = await image_response.arrayBuffer()
+        
+        
+#         return image_data, image_name
+
+# database_url = "https://nyucapstone-7c22c-default-rtdb.firebaseio.com"
+
+# image_data, latest_image_name = await fetch_latest_image_url(database_url)
+
+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
+        # self.imageURL = fetch_latest_image_url("https://nyucapstone-7c22c-default-rtdb.firebaseio.com")
+
+    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 load_image_data(self, image, x, y):
+    #     data = Uint8ClampedArray.new(to_js(self.imageURL.image_url.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)
+    
+    # def draw_image(self, img, x, y):
+    #     self.ctx.drawImage(img, x, y)
+
+    @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,
+        firebase_image_data=None,
+    ) -> 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.np_image = np.array([])
+        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
+        self.firebase_image_data = firebase_image_data
+
+    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 // 4  # Reduce the grid size for more lines
+        x0, y0 = self.view_pos
+        x0 = (-x0) % step
+        y0 = (-y0) % step
+
+        ctx = self.canvas[0].ctx  # Access the CanvasRenderingContext2D object
+
+        ctx.fillStyle = "white"  # Change the fill style to white
+        ctx.fillRect(0, 0, w, h)
+        ctx.strokeStyle = "rgba(0, 0, 0, 0.55)"  # Change the stroke style to transparent black
+        ctx.lineWidth = 0.5  # Make the grid lines thinner
+
+        # Draw horizontal lines
+        for y in range(y0, h + step, step):
+            ctx.beginPath()
+            ctx.moveTo(0, y)
+            ctx.lineTo(w, y)
+            ctx.stroke()
+
+        # Draw vertical lines
+        for x in range(x0, w + step, step):
+            ctx.beginPath()
+            ctx.moveTo(x, 0)
+            ctx.lineTo(x, h)
+            ctx.stroke()
+
+    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,
+        )
+
+
+    async def load_image(self, image_data):
+        # original testing, not being called
+        pil_image = Image.open(io.BytesIO(image_data.to_py()))
+        np_image = np.array(pil_image)
+
+        self.canvas[1].put_image_data(np_image, 0, 0)
+
+
+
+    def draw_buffer(self):
+        self.canvas[1].clear()
+        self.canvas[1].put_image_data(self.buffer, 0, 0)
+        #print(f"self buffer: {self.buffer}") 
+
+        # self.canvas[1].put_image_data(self.firebase_image_data, 0, 0)
+        # print(f"self buffer: {self.firebase_image_data}") 
+
+
+    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,843 @@
+<!DOCTYPE html>
+<html>
+<head>
+<title>Stablediffusion Infinity</title>
+<meta charset="utf-8">
+
+
+<link rel="stylesheet" type="text/css" href="https://cdn.jsdelivr.net/gh/lkwq007/[email protected]/css/w2ui.min.css">
+<script type="text/javascript" src="https://cdn.jsdelivr.net/gh/lkwq007/[email protected]/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/[email protected]/js/fabric.min.js"></script>
+<script defer src="https://cdn.jsdelivr.net/gh/lkwq007/[email protected]/js/toolbar.js"></script>
+<link rel="stylesheet" href="https://pyscript.net/alpha/pyscript.css" />
+<script defer src="https://pyscript.net/alpha/pyscript.js"></script>
+
+
+<script src="https://www.gstatic.com/firebasejs/8.10.0/firebase-app.js"></script>
+<script src="https://www.gstatic.com/firebasejs/8.10.0/firebase-analytics.js"></script>
+<script src="https://www.gstatic.com/firebasejs/8.10.0/firebase-storage.js"></script>
+<script src="https://www.gstatic.com/firebasejs/8.10.0/firebase-database.js"></script>
+
+
+<style>
+html, body {
+  width: 100%;
+  height: 100%;
+  margin: 0;
+  padding: 0;
+  overflow: hidden;
+}
+
+
+#container {
+  position: relative;
+  margin:auto;
+  display: block;
+}
+#container > canvas {
+  position: absolute;
+  top: 0;
+  left: 0;
+}
+.control {
+  display: none;
+}
+#outer_container {
+  width: 100%;
+  height: 100vh;
+  overflow: auto;
+}
+
+#hamburger-menu {
+  position: fixed;
+  top: 10px;
+  right: 10px;
+  width: 50px;
+  height: 50px;
+  background-color: #f1f1f1;
+  border-radius: 50%;
+  display: flex;
+  justify-content: center;
+  align-items: center;
+  cursor: pointer;
+  z-index: 1000;
+  overflow: hidden;
+}
+
+#hamburger-menu::before {
+  content: "";
+  position: absolute;
+  top: -50%;
+  left: -50%;
+  width: 200%;
+  height: 200%;
+  background-image: radial-gradient(circle, #00ff00, #00ffff, #ff00ff, #ff0000, #ffff00, #00ff00);
+  background-size: 300% 300%;
+  animation: gradient-animation 6s linear infinite;
+  z-index: -1;
+}
+
+#hamburger-menu i {
+  font-size: 24px;
+  position: relative;
+  z-index: 1;
+}
+
+.fa-bars {
+  position: relative;
+  display: inline-block;
+  width: 24px;
+  height: 2px;
+  background-color: currentColor;
+  transition: background-color 0.3s ease;
+}
+
+.fa-bars::before,
+.fa-bars::after {
+  content: "";
+  position: absolute;
+  left: 0;
+  width: 100%;
+  height: 2px;
+  background-color: currentColor;
+  transition: transform 0.3s ease, opacity 0.3s ease;
+}
+
+.fa-bars::before {
+  top: -6px;
+}
+
+.fa-bars::after {
+  bottom: -6px;
+}
+
+.open .fa-bars {
+  background-color: transparent;
+}
+
+.open .fa-bars::before {
+  transform: translateY(6px) rotate(45deg);
+}
+
+.open .fa-bars::after {
+  transform: translateY(-6px) rotate(-45deg);
+}
+
+@keyframes gradient-animation {
+  0% {
+    background-position: 0% 50%;
+  }
+  50% {
+    background-position: 100% 50%;
+  }
+  100% {
+    background-position: 0% 50%;
+  }
+}
+#toolbar {
+  display: none;
+}
+
+
+
+.generate-button {
+  background-color: #f1f1f1;
+  border: none;
+  color: #333;
+  padding: 10px 20px;
+  text-align: center;
+  text-decoration: none;
+  display: inline-block;
+  font-size: 16px;
+  margin: 4px 2px;
+  cursor: pointer;
+  border-radius: 4px;
+}
+
+
+
+
+</style>
+
+</head>
+<body>
+<div>
+
+<button type="button" class="control" id="export">Export</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="hamburger-menu">
+    <i class="fa-solid fa-bars"></i>
+  </div>
+  <button type="button"id="outpaint">Outpaint</button>
+<div id = "outer_container">
+<div style="position: relative;">
+    <div id="toolbar" style></div>
+</div> 
+<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>
+</div>
+
+
+<script>
+
+
+
+    alert("starting js");
+
+        function toggleToolbar() {
+        console.log("Hamburger menu button clicked");
+        const toolbar = document.getElementById("toolbar");
+        const hamburgerMenu = document.getElementById("hamburger-menu");
+        if (toolbar.style.display === "none" || toolbar.style.display === "") {
+            toolbar.style.display = "block";
+            hamburgerMenu.classList.add("open");
+        } else {
+            toolbar.style.display = "none";
+            hamburgerMenu.classList.remove("open");
+        }
+        }
+
+        function aws(name, x, y) {
+            return `coming from javascript ${name} ${x} ${y}`;
+        }
+
+    
+
+  
+    const { initializeApp } = firebase;
+   
+    const { getStorage, ref, listAll, getDownloadURL, getMetadata, uploadBytesResumable } = firebase.storage;
+
+
+    const firebaseConfig = {
+        apiKey: "AIzaSyCxG7s_Wg6RAC4AQ5ZpkCgt0XcnSqcwt-A",
+        authDomain: "nyucapstone-7c22c.firebaseapp.com",
+        projectId: "nyucapstone-7c22c",
+        storageBucket: "nyucapstone-7c22c.appspot.com",
+        messagingSenderId: "658619789110",
+        appId: "1:658619789110:web:4eb43edacd4bbfcca74d97",
+        measurementId: "G-NCNE4TC0GC",
+        databaseURL: "https://nyucapstone-7c22c-default-rtdb.firebaseio.com/",
+    };
+
+
+
+    const fireapp = initializeApp(firebaseConfig);
+
+
+    function uploadImageToFirebase(base64_str, time_str) {
+    return new Promise((resolve, reject) => {
+        alert("starting to upload");
+        const atob = (str) => {
+            return window.atob(str);
+        };
+
+        const byteCharacters = atob(base64_str);
+        const byteNumbers = new Uint8Array(byteCharacters.length);
+        for (let i = 0; i < byteCharacters.length; i++) {
+            byteNumbers[i] = byteCharacters.charCodeAt(i);
+        }
+    
+        
+        const analytics = firebase.analytics();
+
+        const byteArray = new Uint8Array(byteNumbers);
+        const blob = new Blob([byteArray], {type: "image/png"});
+    
+
+        const storage = firebase.storage(fireapp);
+        
+      
+        const storageRef = firebase.storage().ref(`images/${time_str}.png`);
+      
+        const uploadTask = storageRef.put(blob);
+
+        alert("sucessful upload to firebae");
+        // Replace the successful upload handler with this:
+        uploadTask.on("state_changed", (snapshot) => {
+            // Handle the progress of the upload
+        }, (error) => {
+            // Handle the error during the upload
+            reject(error);
+        }, async () => {
+            // Handle the successful upload
+            const database = firebase.database();
+            const latestImageRef = database.ref("latestImage");
+            const downloadURL = await storageRef.getDownloadURL();
+            await latestImageRef.set({
+            fileName: `${time_str}.png`,
+            downloadURL: downloadURL
+            });
+            resolve();
+        });
+    });
+}
+    
+    
+document.getElementById("hamburger-menu").addEventListener("click", toggleToolbar);
+  alert("js loaded");
+</script>
+
+<py-env>
+- numpy
+- Pillow
+- micropip:
+    - boto3
+- paths:
+    - ./canvas.py
+</py-env>
+
+
+<py-script>
+from pyodide import to_js, create_proxy
+from PIL import Image
+import io
+import time
+import base64
+from collections import deque
+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,
+    update_undo_redo,
+    alert,
+    uploadImageToFirebase,
+    firebase,
+    aws,
+    fetch
+)
+answer = aws("hello", 1, 2)
+console.log(answer)
+
+#addPhoto("demo")
+
+# async def get_latest_image_from_firebase():
+    # alert("get_latest_image_from_firebase called")
+
+    # try:
+        # database = firebase.database()
+        # alert("try called")
+        # latestImageRef = database.ref("latestImage")
+        # latestImageSnapshot = await latestImageRef.once("value")
+        # latestImageInfo = latestImageSnapshot.val()
+
+      
+        # download_url = latestImageInfo["downloadURL"]
+
+       
+        # with pyodide.open_url(download_url) as f:
+           
+            # img = Image.open(f)
+
+        # print("Downloaded image:", str(img))
+        # return img
+    # except Exception as e:
+        # print("Error while getting the latest image from Firebase:", str(e))
+        # return None
+
+async def fetch_latest_image_url(database_url):
+    console.log("fetch_latest_image called")
+    # different methods to call
+    response = await fetch(f"{database_url}/latestImage.json")
+    console.log(f"response status: {response.status}, status text: {response.statusText}")
+    
+    latest_image_data = await response.json()
+    latest_image_data = latest_image_data.to_py()
+
+    image_url = latest_image_data["downloadURL"]
+    image_name = latest_image_data["fileName"]
+    console.log(f"Latest image URL: {image_url}")
+    console.log(f"Latest image name: {image_name}")
+
+    # Fetch the image data as ArrayBuffer
+    image_response = await fetch(image_url)
+    image_data = await image_response.arrayBuffer()
+    
+    
+    return image_data, image_name
+
+
+from canvas import InfCanvas
+
+
+class History:
+    def __init__(self,maxlen=10):
+        self.idx=-1
+        self.undo_lst=deque([],maxlen=maxlen)
+        self.redo_lst=deque([],maxlen=maxlen)
+        self.state=None
+
+    def undo(self):
+        cur=None
+        if len(self.undo_lst):
+            cur=self.undo_lst.pop()
+            self.redo_lst.appendleft(cur)
+        return cur
+    def redo(self):
+        cur=None
+        if len(self.redo_lst):
+            cur=self.redo_lst.popleft()
+            self.undo_lst.append(cur)
+        return cur
+
+    def check(self):
+        return len(self.undo_lst)>0,len(self.redo_lst)>0
+
+    def append(self,state,update=True):
+        self.redo_lst.clear()
+        self.undo_lst.append(state)
+        if update:
+            update_undo_redo(*self.check())
+
+history = History()
+
+base_lst = [None]
+async def draw_canvas() -> None:
+    alert("draw_canvas called")
+    width=1024
+    height=700
+    canvas=InfCanvas(1024,700)
+    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
+
+    # latest_image = await get_latest_image_from_firebase()
+
+    # if latest_image is not None:
+        # Log the URL of the latest image to the console
+        # console.log(f"Latest image URL: {latest_image.url}")
+        # Request the parent window to display the latest image on the canvas
+        # (commented out to fix the indentation error)
+        # window.parent.postMessage({ type: "displayLatestImageOnCanvas", image: latest_image }, "*")
+    # else:
+        # print("No latest image found in Firebase.")
+
+
+async def draw_canvas_func(event):
+    alert("draw_canvas gradio called")
+    try:
+        app=parent.document.querySelector("gradio-app")
+        if app.shadowRoot:
+            app=app.shadowRoot
+        width=app.querySelector("#canvas_width input").value
+        height=app.querySelector("#canvas_height input").value
+        selection_size=app.querySelector("#selection_size input").value
+    except:
+        width=1024
+        height=768
+        selection_size=384
+    document.querySelector("#container").style.width = f"{width}px"
+    document.querySelector("#container").style.height= f"{height}px"
+
+    database_url = "https://nyucapstone-7c22c-default-rtdb.firebaseio.com"
+    image_data, latest_image_name = await fetch_latest_image_url(database_url)
+    pil_image = Image.open(io.BytesIO(image_data.to_py()))
+
+    np_image = np.array(pil_image)
+
+    canvas=InfCanvas(int(width),int(height),selection_size=int(selection_size),firebase_image_data=np_image)
+
+
+    canvas.setup_mouse()
+    canvas.clear_background()
+    canvas.draw_buffer()
+    canvas.draw_selection_box()
+
+    # await canvas.load_image(image_data)
+
+
+    # Update the canvas buffer with the new image data and redraw the buffer
+    h, w, c = canvas.buffer.shape
+    canvas.sync_to_buffer()
+    canvas.buffer_dirty = True
+
+    h_min = min(h, np_image.shape[0])
+    w_min = min(w, np_image.shape[1])
+
+
+
+    #  mask = np_image[:, :, 3:4].repeat(4, axis=2)
+    #  canvas.buffer[mask > 0] = 0
+    #  canvas.buffer[0:h, 0:w, :] += np_image
+
+    mask = np_image[:h_min, :w_min, 3:4].repeat(4, axis=2)
+    canvas.buffer[:h_min, :w_min][mask > 0] = 0
+    canvas.buffer[:h_min, :w_min] += np_image[:h_min, :w_min]
+    
+    
+    
+    canvas.draw_buffer()
+
+    base_lst[0]=canvas
+
+    alert("made it to end of draw_canvas gradio")
+
+
+import js 
+
+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")
+
+    # The rest of the original export_func code
+    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"{filename}.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()
+    if len(event.data) > 2:
+        history.append(base.save())
+
+    # sending the image to firebase here
+    arr = base.export()
+    base64_str = base.numpy_to_base64(arr)
+    time_str = time.strftime("%Y%m%d_%H%M%S")
+
+    # Call the JavaScript function to upload the image to Firebase storage
+    await js.uploadImageToFirebase(base64_str, time_str)
+    
+
+async def history_undo_func(event):
+    base=base_lst[0]
+    if base.buffer_dirty or len(history.redo_lst)>0:
+        state=history.undo()
+    else:
+        history.undo()
+        state=history.undo()
+    if state is not None:
+        base.load(state)
+    update_undo_redo(*history.check())
+    
+async def history_setup_func(event):
+    base=base_lst[0]
+    history.undo_lst.clear()
+    history.redo_lst.clear()
+    history.append(base.save(),update=False)
+
+async def history_redo_func(event):
+    base=base_lst[0]
+    if len(history.undo_lst)>0:
+        state=history.redo()
+    else:
+        history.redo()
+        state=history.redo()
+    if state is not None:
+        base.load(state)
+    update_undo_redo(*history.check())
+
+
+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]
+    # Retrieve the base64 encoded image string from the #upload_content HTML element
+    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()
+    if len(event.data)>2:
+        history.append(base.save())
+
+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[1]=="redo":
+            await history_redo_func(event)
+        elif event.data[1]=="undo":
+            await history_undo_func(event)
+        elif event.data[1]=="history":
+            await history_setup_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()
+)
+</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]
+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()