From aba446015166261f18d9f2949d3d1def47a56f5c Mon Sep 17 00:00:00 2001 From: ssit Date: Wed, 28 Jun 2023 18:46:35 -0400 Subject: [PATCH] roop working --- .gitignore | 156 ++++++++++++++++++++++++++++++ __init__.py | 17 ++++ gradio.py | 0 install.bat | 13 +++ install.py | 7 +- launch.py | 17 ++++ modules/face_restoration.py | 5 + modules/images.py | 0 modules/processing.py | 13 +++ modules/scripts.py | 13 +++ modules/scripts_postprocessing.py | 0 modules/shared.py | 19 ++++ modules/upscaler.py | 12 +++ nodes.py | 50 ++++++++++ utils.py | 30 ++++++ 15 files changed, 351 insertions(+), 1 deletion(-) create mode 100644 .gitignore create mode 100644 __init__.py create mode 100644 gradio.py create mode 100644 install.bat create mode 100644 launch.py create mode 100644 modules/face_restoration.py create mode 100644 modules/images.py create mode 100644 modules/processing.py create mode 100644 modules/scripts.py create mode 100644 modules/scripts_postprocessing.py create mode 100644 modules/shared.py create mode 100644 modules/upscaler.py create mode 100644 nodes.py create mode 100644 utils.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bfbcdc7 --- /dev/null +++ b/.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 diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..dc2edfd --- /dev/null +++ b/__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] diff --git a/gradio.py b/gradio.py new file mode 100644 index 0000000..e69de29 diff --git a/install.bat b/install.bat new file mode 100644 index 0000000..da87ccd --- /dev/null +++ b/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 diff --git a/install.py b/install.py index 56025dc..08842ab 100644 --- a/install.py +++ b/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: diff --git a/launch.py b/launch.py new file mode 100644 index 0000000..4e3e3ae --- /dev/null +++ b/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(" ")]) diff --git a/modules/face_restoration.py b/modules/face_restoration.py new file mode 100644 index 0000000..00f9723 --- /dev/null +++ b/modules/face_restoration.py @@ -0,0 +1,5 @@ +class FaceRestoration: + pass + +def restore_faces(): + pass diff --git a/modules/images.py b/modules/images.py new file mode 100644 index 0000000..e69de29 diff --git a/modules/processing.py b/modules/processing.py new file mode 100644 index 0000000..d525418 --- /dev/null +++ b/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) diff --git a/modules/scripts.py b/modules/scripts.py new file mode 100644 index 0000000..5eae950 --- /dev/null +++ b/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 diff --git a/modules/scripts_postprocessing.py b/modules/scripts_postprocessing.py new file mode 100644 index 0000000..e69de29 diff --git a/modules/shared.py b/modules/shared.py new file mode 100644 index 0000000..0126386 --- /dev/null +++ b/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 = [] diff --git a/modules/upscaler.py b/modules/upscaler.py new file mode 100644 index 0000000..f6d98cf --- /dev/null +++ b/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() diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..ae115b2 --- /dev/null +++ b/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", +} diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..f3929ea --- /dev/null +++ b/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)