Add initial tools, layout, and config.

.gitignore

@@ -0,0 +1,144 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# IDEs
+/.idea/
+/.vscode/
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+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
+
+# 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/
+
+/wandb/
+wandb/

.streamlit/config.toml

@@ -0,0 +1,3 @@
+[server]
+
+maxUploadSize = 700

README.md

@@ -10,4 +10,30 @@ pinned: false
 license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Lora and Embedding Tools
+
+😻 **Lora and Embedding Tools** is a quick toolbelt to help you manipulate Lora and Text Embedding files. This tool provides several functionalities, including rescaling Lora weights, removing CLIP parameters, converting checkpoint files to safetensors format, and whatever else I decide to add in the future.
+
+## Features
+
+- **Rescale Lora Strength**: Adjust the strength of Lora weights by specifying a new scale factor. Rescales the embedded Alpha scale. (No more 0.6, etc)
+- **Remove CLIP Parameters**: Strip out CLIP parameters from a Lora file. If you have an overbaked or overaggressive Lora file, this can rescue it sometimes, or make it more agnostic for other models.
+- **Convert CKPT to Safetensors**: Convert `.ckpt` files to `.safetensors` format to get that pickle smell out of your weights.
+
+## How to Use
+
+### Rescale Lora Strength
+
+1. Specify the new scale factor first.
+2. Upload a `.safetensors` Lora file, conversion begins immediately.
+3. Download the rescaled weights.
+
+### Remove CLIP Parameters
+
+1. Upload a `.safetensors` Lora file.
+2. Download the file with CLIP parameters removed.
+
+### Convert CKPT to Safetensors
+
+1. Upload a `.ckpt` file (maximum size 700MB).
+2. Download the converted `.safetensors` file.

app.py

@@ -0,0 +1,105 @@
+import io
+
+import safetensors
+import streamlit.file_util
+from safetensors.torch import serialize
+from streamlit.runtime.uploaded_file_manager import UploadedFile
+
+from tools import lora_tools, torch_tools
+
+# https://huggingface.co/docs/hub/spaces-config-reference
+
+streamlit.title("Lora and Embedding Tools")
+
+output_dtype = streamlit.radio("Save Precision", ["float16", "float32", "bfloat16"], index=0)
+streamlit.container()
+col1, col2 = streamlit.columns(2, gap="medium")
+
+# A helper method to wipe a download button once invoked
+def completed_download_callback():
+	ui_filedownload_rescale.empty()
+	ui_filedownload_stripclip.empty()
+	ui_filedownload_ckpt.empty()
+
+with col1:
+	# A tool for rescaling the strength of Lora weights
+	streamlit.html("<h3>Rescale Lora Strength</h3>")
+	ui_fileupload_rescale = streamlit.file_uploader("Upload a safetensors lora", key="fileupload_rescale", type=[".safetensors"])  # type: UploadedFile
+	new_scale_factor = streamlit.number_input("Scale Factor", value=1.0, step=0.01, max_value=100.0, min_value=0.01)
+	
+	# Let's preallocate the download button here so it's in the correct column, we can just add the button later.
+	ui_filedownload_rescale = streamlit.empty()
+
+with col2:
+	# A tool for removing CLIP parameters from a Lora file
+	streamlit.html("<h3>Remove CLIP Parameters</h3>")
+	ui_fileupload_stripclip = streamlit.file_uploader("Upload a safetensors lora", key="fileupload_stripclip", type=[".safetensors"])  # type: UploadedFile
+	
+	# Preallocate download button
+	ui_filedownload_stripclip = streamlit.empty()
+	
+	streamlit.html("<hr>")
+	
+	# A tool for converting a .ckpt file to a .safetensors file
+	streamlit.html("<h3>Convert CKPT to Safetensors (700MB max)</h3>")
+	ui_fileupload_ckpt = streamlit.file_uploader("Upload a .ckpt file", key="fileupload_convertckpt", type=[".ckpt"])  # type: UploadedFile
+	
+	# Preallocate download button
+	ui_filedownload_ckpt = streamlit.empty()
+
+# ! Rescale Lora
+if ui_fileupload_rescale and ui_fileupload_rescale.name is not None:
+	lora_metadata = lora_tools.read_safetensors_metadata(ui_fileupload_rescale)
+	new_weights = lora_tools.rescale_lora_alpha(ui_fileupload_rescale, output_dtype, new_scale_factor)
+	
+	new_lora_data = safetensors.torch.save(new_weights, lora_metadata)
+	
+	lora_file_buffer = io.BytesIO()
+	lora_file_buffer.write(new_lora_data)
+	lora_file_buffer.seek(0)
+	
+	file_name = ui_fileupload_rescale.name.rsplit(".", 1)[0]
+	output_name = f"{file_name}_rescaled.safetensors"
+	
+	ui_fileupload_rescale.close()
+	del ui_fileupload_rescale
+	ui_fileupload_rescale.name = None
+	
+	ui_filedownload_rescale.download_button("Download Rescaled Weights", lora_file_buffer, output_name, type="primary")
+
+# ! Remove CLIP Parameters
+if ui_fileupload_stripclip and ui_fileupload_stripclip.name is not None:
+	lora_metadata = lora_tools.read_safetensors_metadata(ui_fileupload_stripclip)
+	stripped_weights = lora_tools.remove_clip_weights(ui_fileupload_stripclip, output_dtype)
+	
+	stripped_lora_data = safetensors.torch.save(stripped_weights, lora_metadata)
+	
+	lora_file_buffer = io.BytesIO()
+	lora_file_buffer.write(stripped_lora_data)
+	lora_file_buffer.seek(0)
+	
+	file_name = ui_fileupload_stripclip.name.rsplit(".", 1)[0]
+	output_name = f"{file_name}_noclip.safetensors"
+	
+	ui_fileupload_stripclip.close()
+	del ui_fileupload_stripclip
+	
+	ui_filedownload_stripclip.download_button("Download Stripped Weights", lora_file_buffer, output_name, type="primary")
+
+# ! Convert Checkpoint to Safetensors
+if ui_fileupload_ckpt and ui_fileupload_ckpt.name is not None:
+	converted_weights = torch_tools.convert_ckpt_to_safetensors(ui_fileupload_ckpt, output_dtype)
+	
+	converted_lora_data = safetensors.torch.save(converted_weights)
+	
+	lora_file_buffer = io.BytesIO()
+	lora_file_buffer.write(converted_lora_data)
+	lora_file_buffer.seek(0)
+	
+	file_name = ui_fileupload_ckpt.name.rsplit(".", 1)[0]
+	output_name = f"{file_name}.safetensors"
+	
+	ui_fileupload_ckpt.close()
+	del ui_fileupload_ckpt
+	
+	ui_filedownload_ckpt.download_button("Download Converted Weights", lora_file_buffer, output_name, type="primary")

pycharm_runner.py

@@ -0,0 +1,9 @@
+# https://discuss.streamlit.io/t/cannot-debug-streamlit-in-pycharm-2023-3-3/61581/2
+
+try:
+	from streamlit.web import bootstrap
+except ImportError:
+	from streamlit import bootstrap
+
+real_script = 'app.py'
+bootstrap.run(real_script, f'pycharm_runner.py {real_script}', [], {})

requirements.txt

@@ -0,0 +1,3 @@
+safetensors
+streamlit
+torch

tools/__init__.py

@@ -0,0 +1,2 @@
+if __name__ == '__main__':
+	print('__main__ not allowed in modules')

tools/lora_tools.py

@@ -0,0 +1,66 @@
+import io
+import json
+
+import safetensors
+import torch
+from safetensors.torch import serialize
+
+from .torch_tools import get_target_dtype_ref
+
+def read_safetensors_metadata(lora_upload: io.BytesIO) -> dict:
+	# This is a simple file structure, the first 8 bytes are the metadata length.
+	# Read (length) bytes starting from [8] to get the metadata (a json string).
+	lora_upload.seek(0)
+	
+	metadata_length = int.from_bytes(lora_upload.read(8), byteorder='little')
+	
+	lora_upload.seek(8)
+	metadata_raw = lora_upload.read(metadata_length)
+	
+	metadata_raw = metadata_raw.decode("utf-8")
+	metadata_raw = metadata_raw.strip()
+	metadata_dict = json.loads(metadata_raw)
+	
+	# Rewind the buffer to the start, we were just peeking at the metadata.
+	lora_upload.seek(0)
+	
+	return metadata_dict.get('__metadata__', {})
+
+def rescale_lora_alpha(lora_upload: io.BytesIO, output_dtype, target_weight: float = 1.0) -> dict:
+	output_dtype = get_target_dtype_ref(output_dtype)
+	
+	loaded_tensors = safetensors.torch.load(lora_upload.getvalue())
+	
+	initial_tensors = {}
+	for tensor_pair in loaded_tensors.items():
+		key, tensor = tensor_pair
+		initial_tensors[key] = tensor.to(dtype=torch.float32)
+	
+	new_tensors = {}
+	for key, val in initial_tensors.items():
+		if key.endswith(".alpha"):
+			val *= target_weight
+		new_tensors[key] = val.to(dtype=output_dtype)
+	
+	return new_tensors
+
+def remove_clip_weights(lora_upload: io.BytesIO, output_dtype) -> dict:
+	output_dtype = get_target_dtype_ref(output_dtype)
+	
+	loaded_tensors = safetensors.torch.load(lora_upload.getvalue())
+	
+	initial_tensors = {}
+	for tensor_pair in loaded_tensors.items():
+		key, tensor = tensor_pair
+		initial_tensors[key] = tensor.to(dtype=torch.float32)
+	
+	filtered_tensors = {}
+	for key, val in initial_tensors.items():
+		if key.startswith("lora_te1") or key.startswith("lora_te2"):
+			continue
+		filtered_tensors[key] = val.to(dtype=output_dtype)
+	
+	return filtered_tensors
+
+if __name__ == '__main__':
+	print('__main__ not allowed in modules')

tools/torch_tools.py

@@ -0,0 +1,33 @@
+import io
+
+import torch
+
+def get_target_dtype_ref(target_dtype: str) -> torch.dtype:
+	if isinstance(target_dtype, torch.dtype):
+		return target_dtype
+	
+	if target_dtype == "float16":
+		return torch.float16
+	elif target_dtype == "float32":
+		return torch.float32
+	elif target_dtype == "bfloat16":
+		return torch.bfloat16
+	else:
+		raise ValueError(f"Invalid target_dtype: {target_dtype}")
+
+def convert_ckpt_to_safetensors(ckpt_upload: io.BytesIO, target_dtype) -> dict:
+	target_dtype = get_target_dtype_ref(target_dtype)
+	ckpt_data = ckpt_upload.getvalue()
+	
+	# Load the checkpoint
+	checkpoint = torch.load(ckpt_data, map_location="cpu")
+	
+	# Convert the checkpoint to a dictionary of tensors
+	tensor_dict = {}
+	for key, val in checkpoint.items():
+		tensor_dict[key] = val.to(dtype=target_dtype)
+	
+	return tensor_dict
+
+if __name__ == '__main__':
+	print('__main__ not allowed in modules')