backend + frontend

packages.txt

@@ -0,0 +1,32 @@
+python3-opencv
+libopencv-dev
+libopencv-core-dev
+pkg-config
+libopencv-imgcodecs-dev
+libopencv-dev
+libopencv-contrib-dev
+build-essential
+cmake
+git
+pkg-config
+libgtk-3-dev
+libavcodec-dev
+libavformat-dev
+libswscale-dev
+libv4l-dev
+libxvidcore-dev
+libx264-dev
+libjpeg-dev
+libpng-dev
+libtiff-dev
+gfortran
+openexr
+libatlas-base-dev
+python3-dev
+python3-numpy
+libtbb2
+libtbb-dev
+libdc1394-22-dev
+libopenexr-dev
+libgstreamer-plugins-base1.0-dev
+libgstreamer1.0-dev

requirements.txt

@@ -1,8 +1,12 @@
-datasets
-flask
-flask_cors
-flask_expects_json
-requests
-Pillow
-gradio
-Flask-APScheduler
+--extra-index-url https://download.pytorch.org/whl/cu113 
+torch
+huggingface_hub
+diffusers
+transformers
+scikit-image
+pillow
+opencv-python-headless
+fastapi
+uvicorn
+httpx
+gradio

run.py

@@ -1,3 +1,4 @@
-import subprocess
+import os
 
-subprocess.run(["make", "build-all"], shell=False)
+os.system("cd stablediffusion-infinity/PyPatchMatch && make clean && make")
+os.system("cd stablediffusion-infinity && python app.py")

stablediffusion-infinity/.gitignore

@@ -0,0 +1,8 @@
+__pycache__/
+.ipynb_checkpoints/
+build/
+.idea/
+travis.sh
+*.iml
+.token
+libpatchmatch.so

stablediffusion-infinity/.gitmodules

@@ -0,0 +1,3 @@
+[submodule "PyPatchMatch"]
+	path = PyPatchMatch
+	url = https://github.com/vacancy/PyPatchMatch

stablediffusion-infinity/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

stablediffusion-infinity/PyPatchMatch/.gitignore

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

stablediffusion-infinity/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.

stablediffusion-infinity/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.pc) -fPIC
+CXXFLAGS += $(INCLUDE_DIR)
+LDFLAGS = $(shell pkg-config --cflags --libs opencv.pc) -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
+#

stablediffusion-infinity/PyPatchMatch/README.md

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

stablediffusion-infinity/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);
+            }
+        }
+    }
+}
+

stablediffusion-infinity/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;
+};
+

stablediffusion-infinity/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();
+}
+

stablediffusion-infinity/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;
+};
+

stablediffusion-infinity/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));
+}
+

stablediffusion-infinity/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;
+};
+

stablediffusion-infinity/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};
+}
+

stablediffusion-infinity/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" */
+

stablediffusion-infinity/PyPatchMatch/examples/.gitignore

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

stablediffusion-infinity/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;
+}
+

stablediffusion-infinity/PyPatchMatch/examples/cpp_example_run.sh

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

stablediffusion-infinity/PyPatchMatch/examples/py_example.py

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

stablediffusion-infinity/PyPatchMatch/examples/py_example_global_mask.py

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

stablediffusion-infinity/PyPatchMatch/opencv.pc

@@ -0,0 +1,11 @@
+prefix=/usr
+exec_prefix=${prefix}
+includedir=${prefix}/include
+libdir=${exec_prefix}/lib
+
+Name: opencv
+Description: The opencv library
+Version: 2.x.x
+Cflags: -I${includedir}/opencv4
+#Cflags: -I${includedir}/opencv -I${includedir}/opencv2
+Libs: -L${libdir} -lopencv_calib3d -lopencv_imgproc -lopencv_xobjdetect -lopencv_hdf -lopencv_flann -lopencv_core -lopencv_dpm -lopencv_videoio -lopencv_reg  -lopencv_objdetect -lopencv_stitching -lopencv_saliency  -lopencv_features2d -lopencv_text -lopencv_calib3d -lopencv_line_descriptor -lopencv_superres -lopencv_ml -lopencv_viz -lopencv_optflow -lopencv_videostab -lopencv_bioinspired -lopencv_highgui -lopencv_freetype -lopencv_imgcodecs -lopencv_video -lopencv_photo -lopencv_surface_matching -lopencv_rgbd -lopencv_datasets -lopencv_ximgproc -lopencv_plot -lopencv_face -lopencv_stereo -lopencv_aruco -lopencv_phase_unwrapping -lopencv_bgsegm -lopencv_ccalib -lopencv_imgproc -lopencv_shape -lopencv_xphoto -lopencv_structured_light -lopencv_fuzzy

stablediffusion-infinity/PyPatchMatch/patch_match.py

@@ -0,0 +1,191 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+# File   : patch_match.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 01/09/2020
+#
+# Distributed under terms of the MIT license.
+
+import ctypes
+import os.path as osp
+from typing import Optional, Union
+
+import numpy as np
+from PIL import Image
+
+
+__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)
+    ]
+
+
+PMLIB = ctypes.CDLL(osp.join(osp.dirname(__file__), 'libpatchmatch.so'))
+
+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
+

stablediffusion-infinity/app.py

@@ -0,0 +1,234 @@
+import io
+import base64
+import os
+from random import sample
+from sched import scheduler
+
+import uvicorn
+from fastapi import FastAPI, Response
+from fastapi.staticfiles import StaticFiles
+
+import httpx
+from urllib.parse import urljoin
+
+
+import numpy as np
+import torch
+from torch import autocast
+from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
+from PIL import Image
+from PIL import ImageOps
+import gradio as gr
+import base64
+import skimage
+import skimage.measure
+from utils import *
+
+app = FastAPI()
+
+auth_token = os.environ.get("API_TOKEN") or True
+
+WHITES = 66846720
+MASK = Image.open("mask.png")
+try:
+    SAMPLING_MODE = Image.Resampling.LANCZOS
+except Exception as e:
+    SAMPLING_MODE = Image.LANCZOS
+
+
+blocks = gr.Blocks().queue()
+model = {}
+
+
+def get_model():
+    if "text2img" not in model:
+        text2img = StableDiffusionPipeline.from_pretrained(
+            "CompVis/stable-diffusion-v1-4",
+            revision="fp16",
+            torch_dtype=torch.float16,
+            use_auth_token=auth_token,
+        ).to("cuda")
+        inpaint = StableDiffusionInpaintPipeline(
+            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("cuda")
+
+        # lms = LMSDiscreteScheduler(
+        #     beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear")
+
+        # img2img = StableDiffusionImg2ImgPipeline(
+        #     vae=text2img.vae,
+        #     text_encoder=text2img.text_encoder,
+        #     tokenizer=text2img.tokenizer,
+        #     unet=text2img.unet,
+        #     scheduler=lms,
+        #     safety_checker=text2img.safety_checker,
+        #     feature_extractor=text2img.feature_extractor,
+        # ).to("cuda")
+        # try:
+        #     total_memory = torch.cuda.get_device_properties(0).total_memory // (
+        #         1024 ** 3
+        #     )
+        #     if total_memory <= 5:
+        #         inpaint.enable_attention_slicing()
+        # except:
+        #     pass
+        model["text2img"] = text2img
+        model["inpaint"] = inpaint
+        # model["img2img"] = img2img
+
+    return model["text2img"], model["inpaint"]
+    # model["img2img"]
+
+
+get_model()
+
+
+def run_outpaint(
+    input_image,
+    prompt_text,
+    strength,
+    guidance,
+    step,
+    fill_mode,
+):
+    text2img, inpaint = get_model()
+    sel_buffer = np.array(input_image)
+    img = sel_buffer[:, :, 0:3]
+    mask = sel_buffer[:, :, -1]
+    process_size = 512
+
+    mask_sum = mask.sum()
+    if mask_sum >= WHITES:
+        print("inpaiting with fixed Mask")
+        mask = np.array(MASK)[:, :, 0]
+        img, mask = functbl[fill_mode](img, mask)
+        init_image = Image.fromarray(img)
+        mask = 255 - mask
+        mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
+        mask = mask.repeat(8, axis=0).repeat(8, axis=1)
+        mask_image = Image.fromarray(mask)
+
+        # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
+        with autocast("cuda"):
+            images = inpaint(
+                prompt=prompt_text,
+                init_image=init_image.resize(
+                    (process_size, process_size), resample=SAMPLING_MODE
+                ),
+                mask_image=mask_image.resize((process_size, process_size)),
+                strength=strength,
+                num_inference_steps=step,
+                guidance_scale=guidance,
+            )
+    elif mask_sum > 0 and mask_sum < WHITES:
+        print("inpainting")
+        img, mask = functbl[fill_mode](img, mask)
+        init_image = Image.fromarray(img)
+        mask = 255 - mask
+        mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
+        mask = mask.repeat(8, axis=0).repeat(8, axis=1)
+        mask_image = Image.fromarray(mask)
+
+        # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
+        with autocast("cuda"):
+            images = inpaint(
+                prompt=prompt_text,
+                init_image=init_image.resize(
+                    (process_size, process_size), resample=SAMPLING_MODE
+                ),
+                mask_image=mask_image.resize((process_size, process_size)),
+                strength=strength,
+                num_inference_steps=step,
+                guidance_scale=guidance,
+            )
+    else:
+        print("text2image")
+        with autocast("cuda"):
+            images = text2img(
+                prompt=prompt_text, height=process_size, width=process_size,
+            )
+
+    return images['sample'][0], images["nsfw_content_detected"][0]
+
+
+with blocks as demo:
+
+    with gr.Row():
+
+        with gr.Column(scale=3, min_width=270):
+            sd_prompt = gr.Textbox(
+                label="Prompt", placeholder="input your prompt here", lines=4
+            )
+        with gr.Column(scale=2, min_width=150):
+            sd_strength = gr.Slider(
+                label="Strength", minimum=0.0, maximum=1.0, value=0.75, step=0.01
+            )
+        with gr.Column(scale=1, min_width=150):
+            sd_step = gr.Number(label="Step", value=50, precision=0)
+            sd_guidance = gr.Number(label="Guidance", value=7.5)
+    with gr.Row():
+        with gr.Column(scale=4, min_width=600):
+            init_mode = gr.Radio(
+                label="Init mode",
+                choices=[
+                    "patchmatch",
+                    "edge_pad",
+                    "cv2_ns",
+                    "cv2_telea",
+                    "gaussian",
+                    "perlin",
+                ],
+                value="patchmatch",
+                type="value",
+            )
+
+    model_input = gr.Image(label="Input", type="pil", image_mode="RGBA")
+    proceed_button = gr.Button("Proceed", elem_id="proceed")
+    model_output = gr.Image(label="Output")
+    is_nsfw = gr.JSON()
+
+    proceed_button.click(
+        fn=run_outpaint,
+        inputs=[
+            model_input,
+            sd_prompt,
+            sd_strength,
+            sd_guidance,
+            sd_step,
+            init_mode,
+        ],
+        outputs=[model_output, is_nsfw],
+    )
+
+
+blocks.config['dev_mode'] = False
+
+S3_HOST = "https://s3.amazonaws.com"
+
+
+@app.get("/uploads/{path:path}")
+async def uploads(path: str, response: Response):
+    async with httpx.AsyncClient() as client:
+        proxy = await client.get(f"{S3_HOST}/{path}")
+    response.body = proxy.content
+    response.status_code = proxy.status_code
+    response.headers['Access-Control-Allow-Origin'] = '*'
+    response.headers['Access-Control-Allow-Methods'] = 'POST, GET, DELETE, OPTIONS'
+    response.headers['Access-Control-Allow-Headers'] = 'Authorization, Content-Type'
+    return response
+
+
+app = gr.mount_gradio_app(app, blocks, "/gradio",
+                          gradio_api_url="http://localhost:7860/gradio")
+
+app.mount("/", StaticFiles(directory="../static", html=True), name="static")
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860,
+                log_level="debug", reload=False)

stablediffusion-infinity/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
+
+
+##########

stablediffusion-infinity/readme.md

@@ -0,0 +1,93 @@
+# stablediffusion-infinity
+
+Outpainting with Stable Diffusion on an infinite canvas.
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/lkwq007/stablediffusion-infinity/blob/master/stablediffusion_infinity_colab.ipynb)
+
+Start with init_image:
+
+
+
+https://user-images.githubusercontent.com/1665437/190231611-fc263115-0fb9-4f2d-a71b-7e500c1e311d.mp4
+
+
+Start with text2img:
+
+https://user-images.githubusercontent.com/1665437/190212025-f4a82c46-0ff1-4ca2-b79b-6c81601e3eed.mp4
+
+
+It is recommended to run the notebook on a local server for better interactive control. 
+
+The notebook might work on Windows (see this issue https://github.com/lkwq007/stablediffusion-infinity/issues/12 for more information) and Apple Silicon devices (untested, check guide here: https://huggingface.co/docs/diffusers/optimization/mps). 
+
+## Status
+
+This project mainly works as a proof of concept. In that case, the UI design is relatively weak, and the quality of results is not guaranteed. 
+You may need to do prompt engineering or change the size of the selection box to get better outpainting results. 
+
+Pull requests are welcome for better UI control, ideas to achieve better results, or any other improvements. 
+
+## Setup environment
+setup with `environment.yml`
+```
+git clone --recurse-submodules https://github.com/lkwq007/stablediffusion-infinity
+cd stablediffusion-infinity
+conda env create -f environment.yml
+```
+
+if the `environment.yml` doesn't work for you, you may install dependencies manually: 
+```
+conda create -n sd-inf python=3.10
+conda activate sd-inf
+conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
+conda install scipy scikit-image
+conda install -c conda-forge diffusers transformers ftfy
+pip install opencv-python
+pip install gradio==3.4
+```
+
+For windows, you may need to replace `pip install opencv-python` with `conda install -c conda-forge opencv`
+## CPP library (optional)
+
+Note that `opencv` library (e.g. `libopencv-dev`/`opencv-devel`, the package name may differ on different distributions) is required for `PyPatchMatch`. You may need to install `opencv` by yourself. If no `opencv` installed, the `patch_match` option (usually better quality) won't work. 
+
+## How-to
+
+```
+conda activate sd-inf
+python app.py
+```
+
+## Running with Docker
+
+This should get you started without needing to manually install anything, except for having an environment with Docker installed and an Nvidia GPU.
+This has been tested on Docker Desktop on Windows 10 using the WSL2 backend.
+
+First, update the .env file with your Huggingface token from https://huggingface.co/settings/tokens
+
+Open your shell that has docker and run these commands
+
+```
+cd stablediffusion-infinity
+docker-compose build
+docker-compose up
+```
+
+Watch the log for the url to open in your browser. Choose the one that starts with http://127.0.0.1:8888/
+
+
+## FAQs
+
+- Troubleshooting on Windows: 
+  - https://github.com/lkwq007/stablediffusion-infinity/issues/12
+- False positive rate of safety checker is quite high: 
+  - https://github.com/lkwq007/stablediffusion-infinity/issues/8#issuecomment-1248448453
+- What is the init_mode
+  - init_mode indicates how to fill the empty/masked region, usually `patch_match` is better than others
+- The GUI is lagging on colab
+  - It is recommended to run the notebook on a local server since the interactions and canvas content updates are actually handled by the python backend on the serverside, and that's how `ipycanvas` works
+  - colab doesn't support the latest version of `ipycanvas`, which may have better performance
+
+## Credit
+
+The code of `perlin2d.py` is from https://stackoverflow.com/questions/42147776/producing-2d-perlin-noise-with-numpy/42154921#42154921 and is **not** included in the scope of LICENSE used in this repo.

stablediffusion-infinity/utils.py

@@ -0,0 +1,140 @@
+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
+
+from PyPatchMatch import patch_match
+
+
+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
+
+
+functbl = {
+    "gaussian": gaussian_noise,
+    "perlin": perlin_noise,
+    "edge_pad": edge_pad,
+    "patchmatch": patch_match_func if (os.name != "nt" and patch_match_compiled) else edge_pad,
+    "cv2_ns": cv2_ns,
+    "cv2_telea": cv2_telea,
+    "mean_fill": mean_fill,
+}