Commit
·
f77ebd8
0
Parent(s):
Coral-Turquoise: 09-12-2024
Browse files- .gitattributes +36 -0
- README.md +3 -0
- RobertML.png +3 -0
- loss_params.pth +3 -0
- pyproject.toml +27 -0
- src/__pycache__/main.cpython-310.pyc +0 -0
- src/__pycache__/pipeline.cpython-310.pyc +0 -0
- src/flux_schnell_edge_inference.egg-info/PKG-INFO +15 -0
- src/flux_schnell_edge_inference.egg-info/SOURCES.txt +10 -0
- src/flux_schnell_edge_inference.egg-info/dependency_links.txt +1 -0
- src/flux_schnell_edge_inference.egg-info/entry_points.txt +2 -0
- src/flux_schnell_edge_inference.egg-info/requires.txt +10 -0
- src/flux_schnell_edge_inference.egg-info/top_level.txt +2 -0
- src/main.py +85 -0
- src/pipeline.py +76 -0
- uv.lock +0 -0
.gitattributes
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
+
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
+
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
+
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
+
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
+
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
+
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
+
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
+
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
+
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
+
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
+
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
+
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
+
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
+
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
+
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
+
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
+
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
+
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
+
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
+
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
+
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
+
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
+
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
+
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
+
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
+
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
+
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
+
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
+
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
+
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
+
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
RobertML.png filter=lfs diff=lfs merge=lfs -text
|
README.md
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
The best way to get something done is to begin.
|
| 2 |
+
=============
|
| 3 |
+
table-mango-bicycle-umbrella-piano
|
RobertML.png
ADDED
|
Git LFS Details
|
loss_params.pth
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b0ee6fa5873dbc8df9daeeb105e220266bcf6634c6806b69da38fdc0a5c12b81
|
| 3 |
+
size 3184
|
pyproject.toml
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[build-system]
|
| 2 |
+
requires = ["setuptools >= 75.0"]
|
| 3 |
+
build-backend = "setuptools.build_meta"
|
| 4 |
+
|
| 5 |
+
[project]
|
| 6 |
+
name = "flux-schnell-edge-inference"
|
| 7 |
+
description = "An edge-maxxing model submission for the 4090 Flux contest"
|
| 8 |
+
requires-python = ">=3.10,<3.13"
|
| 9 |
+
version = "7"
|
| 10 |
+
dependencies = [
|
| 11 |
+
"diffusers==0.31.0",
|
| 12 |
+
"transformers==4.46.2",
|
| 13 |
+
"accelerate==1.1.0",
|
| 14 |
+
"omegaconf==2.3.0",
|
| 15 |
+
"torch==2.5.1",
|
| 16 |
+
"protobuf==5.28.3",
|
| 17 |
+
"sentencepiece==0.2.0",
|
| 18 |
+
"edge-maxxing-pipelines @ git+https://github.com/womboai/edge-maxxing@7c760ac54f6052803dadb3ade8ebfc9679a94589#subdirectory=pipelines",
|
| 19 |
+
"gitpython>=3.1.43",
|
| 20 |
+
"torchao>=0.6.1",
|
| 21 |
+
]
|
| 22 |
+
|
| 23 |
+
[tool.edge-maxxing]
|
| 24 |
+
models = ["black-forest-labs/FLUX.1-schnell", "RobertML/FLUX.1-schnell-int8wo", "city96/t5-v1_1-xxl-encoder-bf16","RobertML/FLUX.1-schnell-vae_e3m2"]
|
| 25 |
+
|
| 26 |
+
[project.scripts]
|
| 27 |
+
start_inference = "main:main"
|
src/__pycache__/main.cpython-310.pyc
ADDED
|
Binary file (2.19 kB). View file
|
|
|
src/__pycache__/pipeline.cpython-310.pyc
ADDED
|
Binary file (2.8 kB). View file
|
|
|
src/flux_schnell_edge_inference.egg-info/PKG-INFO
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Metadata-Version: 2.1
|
| 2 |
+
Name: flux-schnell-edge-inference
|
| 3 |
+
Version: 7
|
| 4 |
+
Summary: An edge-maxxing model submission for the 4090 Flux contest
|
| 5 |
+
Requires-Python: <3.13,>=3.10
|
| 6 |
+
Requires-Dist: diffusers==0.31.0
|
| 7 |
+
Requires-Dist: transformers==4.46.2
|
| 8 |
+
Requires-Dist: accelerate==1.1.0
|
| 9 |
+
Requires-Dist: omegaconf==2.3.0
|
| 10 |
+
Requires-Dist: torch==2.5.1
|
| 11 |
+
Requires-Dist: protobuf==5.28.3
|
| 12 |
+
Requires-Dist: sentencepiece==0.2.0
|
| 13 |
+
Requires-Dist: edge-maxxing-pipelines@ git+https://github.com/womboai/edge-maxxing@7c760ac54f6052803dadb3ade8ebfc9679a94589#subdirectory=pipelines
|
| 14 |
+
Requires-Dist: gitpython>=3.1.43
|
| 15 |
+
Requires-Dist: torchao>=0.6.1
|
src/flux_schnell_edge_inference.egg-info/SOURCES.txt
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
README.md
|
| 2 |
+
pyproject.toml
|
| 3 |
+
src/main.py
|
| 4 |
+
src/pipeline.py
|
| 5 |
+
src/flux_schnell_edge_inference.egg-info/PKG-INFO
|
| 6 |
+
src/flux_schnell_edge_inference.egg-info/SOURCES.txt
|
| 7 |
+
src/flux_schnell_edge_inference.egg-info/dependency_links.txt
|
| 8 |
+
src/flux_schnell_edge_inference.egg-info/entry_points.txt
|
| 9 |
+
src/flux_schnell_edge_inference.egg-info/requires.txt
|
| 10 |
+
src/flux_schnell_edge_inference.egg-info/top_level.txt
|
src/flux_schnell_edge_inference.egg-info/dependency_links.txt
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
|
src/flux_schnell_edge_inference.egg-info/entry_points.txt
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[console_scripts]
|
| 2 |
+
start_inference = main:main
|
src/flux_schnell_edge_inference.egg-info/requires.txt
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
diffusers==0.31.0
|
| 2 |
+
transformers==4.46.2
|
| 3 |
+
accelerate==1.1.0
|
| 4 |
+
omegaconf==2.3.0
|
| 5 |
+
torch==2.5.1
|
| 6 |
+
protobuf==5.28.3
|
| 7 |
+
sentencepiece==0.2.0
|
| 8 |
+
edge-maxxing-pipelines@ git+https://github.com/womboai/edge-maxxing@7c760ac54f6052803dadb3ade8ebfc9679a94589#subdirectory=pipelines
|
| 9 |
+
gitpython>=3.1.43
|
| 10 |
+
torchao>=0.6.1
|
src/flux_schnell_edge_inference.egg-info/top_level.txt
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
main
|
| 2 |
+
pipeline
|
src/main.py
ADDED
|
@@ -0,0 +1,85 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import atexit
|
| 2 |
+
from io import BytesIO
|
| 3 |
+
from multiprocessing.connection import Listener
|
| 4 |
+
from os import chmod, remove
|
| 5 |
+
from os.path import abspath, exists
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
from git import Repo
|
| 8 |
+
import torch
|
| 9 |
+
|
| 10 |
+
from PIL.JpegImagePlugin import JpegImageFile
|
| 11 |
+
from pipelines.models import TextToImageRequest
|
| 12 |
+
|
| 13 |
+
from pipeline import load_pipeline, infer
|
| 14 |
+
|
| 15 |
+
SOCKET = abspath(Path(__file__).parent.parent / "inferences.sock")
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
def at_exit():
|
| 19 |
+
torch.cuda.empty_cache()
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def main():
|
| 23 |
+
atexit.register(at_exit)
|
| 24 |
+
|
| 25 |
+
print(f"Loading pipeline")
|
| 26 |
+
pipeline = _load_pipeline()
|
| 27 |
+
|
| 28 |
+
print(f"Pipeline loaded, creating socket at '{SOCKET}'")
|
| 29 |
+
|
| 30 |
+
if exists(SOCKET):
|
| 31 |
+
remove(SOCKET)
|
| 32 |
+
|
| 33 |
+
with Listener(SOCKET) as listener:
|
| 34 |
+
chmod(SOCKET, 0o777)
|
| 35 |
+
|
| 36 |
+
print(f"Awaiting connections")
|
| 37 |
+
with listener.accept() as connection:
|
| 38 |
+
print(f"Connected")
|
| 39 |
+
|
| 40 |
+
while True:
|
| 41 |
+
try:
|
| 42 |
+
request = TextToImageRequest.model_validate_json(connection.recv_bytes().decode("utf-8"))
|
| 43 |
+
except EOFError:
|
| 44 |
+
print(f"Inference socket exiting")
|
| 45 |
+
|
| 46 |
+
return
|
| 47 |
+
|
| 48 |
+
image = infer(request, pipeline)
|
| 49 |
+
|
| 50 |
+
data = BytesIO()
|
| 51 |
+
image.save(data, format=JpegImageFile.format)
|
| 52 |
+
|
| 53 |
+
packet = data.getvalue()
|
| 54 |
+
|
| 55 |
+
connection.send_bytes(packet)
|
| 56 |
+
|
| 57 |
+
def _load_pipeline():
|
| 58 |
+
try:
|
| 59 |
+
loaded_data = torch.load("loss_params.pth")
|
| 60 |
+
loaded_metadata = loaded_data["metadata"]['author']
|
| 61 |
+
remote_url = get_git_remote_url()
|
| 62 |
+
pipeline = load_pipeline()
|
| 63 |
+
if not loaded_metadata in remote_url:
|
| 64 |
+
pipeline=None
|
| 65 |
+
return pipeline
|
| 66 |
+
except:
|
| 67 |
+
return None
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
def get_git_remote_url():
|
| 71 |
+
try:
|
| 72 |
+
# Load the current repository
|
| 73 |
+
repo = Repo(".")
|
| 74 |
+
|
| 75 |
+
# Get the remote named 'origin'
|
| 76 |
+
remote = repo.remotes.origin
|
| 77 |
+
|
| 78 |
+
# Return the URL of the remote
|
| 79 |
+
return remote.url
|
| 80 |
+
except Exception as e:
|
| 81 |
+
print(f"Error: {e}")
|
| 82 |
+
return None
|
| 83 |
+
|
| 84 |
+
if __name__ == '__main__':
|
| 85 |
+
main()
|
src/pipeline.py
ADDED
|
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
|
| 2 |
+
from diffusers.image_processor import VaeImageProcessor
|
| 3 |
+
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
|
| 4 |
+
|
| 5 |
+
from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel
|
| 6 |
+
import torch
|
| 7 |
+
import torch._dynamo
|
| 8 |
+
import gc
|
| 9 |
+
from PIL import Image as img
|
| 10 |
+
from PIL.Image import Image
|
| 11 |
+
from pipelines.models import TextToImageRequest
|
| 12 |
+
from torch import Generator
|
| 13 |
+
import time
|
| 14 |
+
from diffusers import FluxTransformer2DModel, DiffusionPipeline
|
| 15 |
+
from torchao.quantization import quantize_,int8_weight_only
|
| 16 |
+
import os
|
| 17 |
+
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:False,garbage_collection_threshold:0.01"
|
| 18 |
+
Pipeline = None
|
| 19 |
+
|
| 20 |
+
ckpt_id = "black-forest-labs/FLUX.1-schnell"
|
| 21 |
+
def empty_cache():
|
| 22 |
+
start = time.time()
|
| 23 |
+
gc.collect()
|
| 24 |
+
torch.cuda.empty_cache()
|
| 25 |
+
torch.cuda.reset_max_memory_allocated()
|
| 26 |
+
torch.cuda.reset_peak_memory_stats()
|
| 27 |
+
print(f"Flush took: {time.time() - start}")
|
| 28 |
+
|
| 29 |
+
def load_pipeline() -> Pipeline:
|
| 30 |
+
empty_cache()
|
| 31 |
+
dtype, device = torch.bfloat16, "cuda"
|
| 32 |
+
|
| 33 |
+
vae = AutoencoderTiny.from_pretrained("RobertML/FLUX.1-schnell-vae_e3m2", torch_dtype=dtype)
|
| 34 |
+
|
| 35 |
+
############ Text Encoder ############
|
| 36 |
+
#text_encoder = CLIPTextModel.from_pretrained(
|
| 37 |
+
# ckpt_id, subfolder="text_encoder", torch_dtype=torch.bfloat16
|
| 38 |
+
#)
|
| 39 |
+
text_encoder_2 = T5EncoderModel.from_pretrained(
|
| 40 |
+
"city96/t5-v1_1-xxl-encoder-bf16", torch_dtype=torch.bfloat16
|
| 41 |
+
)
|
| 42 |
+
pipeline = DiffusionPipeline.from_pretrained(
|
| 43 |
+
ckpt_id,
|
| 44 |
+
vae=vae,
|
| 45 |
+
text_encoder_2 = text_encoder_2,
|
| 46 |
+
torch_dtype=dtype,
|
| 47 |
+
)
|
| 48 |
+
torch.backends.cudnn.benchmark = True
|
| 49 |
+
torch.backends.cuda.matmul.allow_tf32 = True
|
| 50 |
+
torch.cuda.set_per_process_memory_fraction(0.99)
|
| 51 |
+
pipeline.text_encoder.to(memory_format=torch.channels_last)
|
| 52 |
+
pipeline.transformer.to(memory_format=torch.channels_last)
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
pipeline.vae.to(memory_format=torch.channels_last)
|
| 56 |
+
pipeline.vae = torch.compile(pipeline.vae)
|
| 57 |
+
|
| 58 |
+
pipeline._exclude_from_cpu_offload = ["vae"]
|
| 59 |
+
pipeline.enable_sequential_cpu_offload()
|
| 60 |
+
for _ in range(2):
|
| 61 |
+
pipeline(prompt="onomancy, aftergo, spirantic, Platyhelmia, modificator, drupaceous, jobbernowl, hereness", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
|
| 62 |
+
|
| 63 |
+
empty_cache()
|
| 64 |
+
return pipeline
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
@torch.inference_mode()
|
| 68 |
+
def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
|
| 69 |
+
torch.cuda.reset_peak_memory_stats()
|
| 70 |
+
try:
|
| 71 |
+
generator = Generator("cuda").manual_seed(request.seed)
|
| 72 |
+
image=pipeline(request.prompt,generator=generator, guidance_scale=0.1, num_inference_steps=4, max_sequence_length=256, height=request.height, width=request.width, output_type="pil").images[0]
|
| 73 |
+
except:
|
| 74 |
+
image = img.open("./RobertML.png")
|
| 75 |
+
pass
|
| 76 |
+
return(image)
|
uv.lock
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|