roop working

.gitignore

@@ -0,0 +1,156 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintainted in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Other
+*.ipynb
+*.onnx

__init__.py

@@ -0,0 +1,17 @@
+import sys
+import os
+repo_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.insert(0, repo_dir)
+modules = sys.modules.copy()
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+
+# Clean up imports
+sys.path.remove(repo_dir)
+modules_to_remove = []
+for module in sys.modules:
+    if module not in modules:
+        modules_to_remove.append(module)
+for module in modules_to_remove:
+    del sys.modules[module]

install.bat

@@ -0,0 +1,13 @@
+@echo off
+:: Exit if embedded python is not found
+if not exist ..\..\..\python_embeded\python.exe (
+    echo Embedded python not found. Please install manually.
+    pause
+    exit /b 1
+)
+
+:: Install the package
+echo Installing roop requirements...
+..\..\..\python_embeded\python.exe install.py
+echo Finished installing roop requirements.
+pause

install.py

@@ -1,9 +1,10 @@
-import launch
 import os
 import pkg_resources
 import sys
 from tqdm import tqdm
 import urllib.request
+sys.path.append(os.path.dirname(os.path.realpath(__file__)))
+import launch
 
 req_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), "requirements.txt")
 
@@ -24,6 +25,10 @@ if not os.path.exists(models_dir):
 if not os.path.exists(model_path):
     download(model_url, model_path)
 
+# Copy model to ./scripts/ using a hard link
+dst = os.path.join(os.path.dirname(os.path.realpath(__file__)), "scripts", model_name)
+os.link(model_path, dst)
+
 print("Checking roop requirements")
 with open(req_file) as file:
     for package in file:

launch.py

@@ -0,0 +1,17 @@
+import importlib.util
+import subprocess
+import sys
+
+
+def is_installed(package):
+    try:
+        spec = importlib.util.find_spec(package)
+    except ModuleNotFoundError:
+        return False
+
+    return spec is not None
+
+
+def run_pip(command, desc):
+    python = sys.executable
+    subprocess.check_call([python, "-m", "pip", *command.split(" ")])

modules/face_restoration.py

@@ -0,0 +1,5 @@
+class FaceRestoration:
+    pass
+
+def restore_faces():
+    pass

modules/processing.py

@@ -0,0 +1,13 @@
+class StableDiffusionProcessing:
+
+    def __init__(self, init_imgs):
+        self.init_images = init_imgs
+        self.width = init_imgs[0].width
+        self.height = init_imgs[0].height
+        self.extra_generation_params = {}
+
+
+class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
+
+    def __init__(self, init_img):
+        super().__init__(init_img)

modules/scripts.py

@@ -0,0 +1,13 @@
+import os
+
+
+class Script:
+    pass
+
+
+def basedir():
+    return os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+class PostprocessImageArgs:
+    pass

modules/shared.py

@@ -0,0 +1,19 @@
+class Options:
+    img2img_background_color = "#ffffff"  # Set to white for now
+
+
+class State:
+    interrupted = False
+
+    def begin(self):
+        pass
+
+    def end(self):
+        pass
+
+
+opts = Options()
+state = State()
+cmd_opts = None
+sd_upscalers = []
+face_restorers = []

modules/upscaler.py

@@ -0,0 +1,12 @@
+class Upscaler:
+
+    def upscale(self, img, scale, selected_model: str = None):
+        pass
+
+
+class UpscalerData:
+    name = ""
+    data_path = ""
+
+    def __init__(self):
+        self.scaler = Upscaler()

nodes.py

@@ -0,0 +1,50 @@
+import os
+from modules.processing import StableDiffusionProcessingImg2Img
+from scripts.faceswap import FaceSwapScript, get_models
+from utils import batch_tensor_to_pil, batched_pil_to_tensor, tensor_to_pil
+
+
+def model_names():
+    models = get_models()
+    return {os.path.basename(x): x for x in models}
+
+
+class roop:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "reference_image": ("IMAGE",),
+                "swap_model": (list(model_names().keys()),),
+                # Comma separated face number(s)
+                "faces_index": ("STRING", {"default": "0"}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "execute"
+    CATEGORY = "image/postprocessing"
+
+    def execute(self, image, reference_image, swap_model, faces_index):
+        script = FaceSwapScript()
+        pil_images = batch_tensor_to_pil(image)
+        source = tensor_to_pil(reference_image)
+        p = StableDiffusionProcessingImg2Img(pil_images)
+        script.process(
+            p=p, img=source, enable=True, faces_index=faces_index, model=swap_model, 
+            face_restorer_name=None, face_restorer_visibility=None, 
+            upscaler_name=None, upscaler_scale=None, upscaler_visibility=None, 
+            swap_in_source=True, swap_in_generated=True
+        )
+        result = batched_pil_to_tensor(p.init_images)
+        return (result,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "roop": roop,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "roop": "roop",
+}

utils.py

@@ -0,0 +1,30 @@
+from PIL import Image
+import numpy as np
+import torch
+
+
+def tensor_to_pil(img_tensor, batch_index=0):
+    # Convert tensor of shape [batch_size, channels, height, width] at the batch_index to PIL Image
+    img_tensor = img_tensor[batch_index].unsqueeze(0)
+    i = 255. * img_tensor.cpu().numpy()
+    img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8).squeeze())
+    return img
+
+
+def batch_tensor_to_pil(img_tensor):
+    # Convert tensor of shape [batch_size, channels, height, width] to a list of PIL Images
+    return [tensor_to_pil(img_tensor, i) for i in range(img_tensor.shape[0])]
+
+
+def pil_to_tensor(image):
+    # Takes a PIL image and returns a tensor of shape [1, height, width, channels]
+    image = np.array(image).astype(np.float32) / 255.0
+    image = torch.from_numpy(image).unsqueeze(0)
+    if len(image.shape) == 3:  # If the image is grayscale, add a channel dimension
+        image = image.unsqueeze(-1)
+    return image
+
+
+def batched_pil_to_tensor(images):
+    # Takes a list of PIL images and returns a tensor of shape [batch_size, height, width, channels]
+    return torch.cat([pil_to_tensor(image) for image in images], dim=0)