Skip to content

Commit

Permalink
fix: move parser outof model loader
Browse files Browse the repository at this point in the history
  • Loading branch information
Mikubill committed Feb 13, 2023
1 parent 1da97f4 commit 1a2d545
Show file tree
Hide file tree
Showing 2 changed files with 188 additions and 28 deletions.
160 changes: 160 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,160 @@
# 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

# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#pdm.lock
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
# in version control.
# https://pdm.fming.dev/#use-with-ide
.pdm.toml

# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__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 maintained 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/
56 changes: 28 additions & 28 deletions scripts/controlnet.py
Original file line number Diff line number Diff line change
Expand Up @@ -174,18 +174,17 @@ def refresh_all_models(*dropdowns):
ctrls += (refresh_models, )

def create_canvas(h, w):
return np.zeros(shape=(h, w, 3), dtype=np.uint8) + 255, True
return np.zeros(shape=(h, w, 3), dtype=np.uint8) + 255

canvas_state = gr.State(False)
canvas_width = gr.Slider(label="Canvas Width", minimum=256, maximum=1024, value=512, step=1)
canvas_height = gr.Slider(label="Canvas Height", minimum=256, maximum=1024, value=512, step=1)
create_button = gr.Button(label="Start", value='Open drawing canvas!')
input_image = gr.Image(source='upload', type='numpy', tool='sketch')
gr.Markdown(value='Do not forget to change your brush width to make it thinner. (Gradio do not allow developers to set brush width so you need to do it manually.) '
'Just click on the small pencil icon in the upper right corner of the above block.')

create_button.click(fn=create_canvas, inputs=[canvas_width, canvas_height], outputs=[input_image, canvas_state])
ctrls += (canvas_width, canvas_height, create_button, input_image, canvas_state, scribble_mode)
create_button.click(fn=create_canvas, inputs=[canvas_width, canvas_height], outputs=[input_image])
ctrls += (canvas_width, canvas_height, create_button, input_image, scribble_mode)

return ctrls

Expand Down Expand Up @@ -216,7 +215,9 @@ def restore_networks():
self.latest_network = None

enabled, module, model, weight, _ = args[:5]
_, _, _, image, canvas_state, scribble_mode = args[5:]
_, _, _, image, scribble_mode = args[5:]

print("called here")

if not enabled:
restore_networks()
Expand All @@ -241,43 +242,42 @@ def restore_networks():
if not os.path.exists(model_path):
raise ValueError(f"file not found: {model_path}")

print(f"using preprocessor: {module}, model: {model}")
print(f"loading preprocessor: {module}, model: {model}")
network = PlugableControlModel(model_path, os.path.join(cn_models_dir, "cldm_v15.yaml"), weight)
network.to(p.sd_model.device, dtype=p.sd_model.dtype)
network.hook(unet)

print(f"ControlNet model {model} loaded.")
self.latest_network = network

input_image = HWC3(image['image'])
if canvas_state:
print("using mask as input")
input_image = HWC3(image['mask'][:, :, 0])
input_image = HWC3(image['image'])
if 255 - np.mean(input_image) < 5:
print("using mask as input")
input_image = HWC3(image['mask'][:, :, 0])

if scribble_mode:
detected_map = np.zeros_like(input_image, dtype=np.uint8)
detected_map[np.min(input_image, axis=2) < 127] = 255
input_image = detected_map
if scribble_mode:
detected_map = np.zeros_like(input_image, dtype=np.uint8)
detected_map[np.min(input_image, axis=2) < 127] = 255
input_image = detected_map

preprocessor = self.preprocessor[self.latest_params[0]]
h, w, bsz = p.height, p.width, p.batch_size
detected_map = preprocessor(input_image)
detected_map = HWC3(detected_map)

control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
control = rearrange(control, 'h w c -> c h w')
control = Resize(h if h>w else w, interpolation=InterpolationMode.BICUBIC)(control)
control = CenterCrop((h, w))(control)
print(control)
preprocessor = self.preprocessor[self.latest_params[0]]
h, w, bsz = p.height, p.width, p.batch_size
detected_map = preprocessor(input_image)
detected_map = HWC3(detected_map)

control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
control = rearrange(control, 'h w c -> c h w')
control = Resize(h if h>w else w, interpolation=InterpolationMode.BICUBIC)(control)
control = CenterCrop((h, w))(control)

self.control = control
control = torch.stack([control for _ in range(bsz)], dim=0)
self.latest_network.notify(control)
self.control = control
control = torch.stack([control for _ in range(bsz)], dim=0)
self.latest_network.notify(control)

self.set_infotext_fields(p, self.latest_params)

def postprocess(self, p, processed, *args):
processed.images.append(ToPILImage()((self.control).clip(0, 255)))
# processed.images.append(ToPILImage()((self.control).clip(0, 255)))
pass

def update_script_args(p, value, arg_idx):
Expand Down

0 comments on commit 1a2d545

Please sign in to comment.