[BUG] IndexError with tensor_sum and rebasin #113
Description
When running with
merge_mode: tensor_sum
rebasin: TrueI get a IndexError: tuple index out of range at "Rebasin iteration 1".
It doesn't happen when rebasin is set to False or when using weigthed_sum instead of tensor_sum.
Likely related to #76.
Output:
2024-03-11 10:59:37 INFO [root] Loading: ~/Downloads/stable-diffusion-webui/models/Stable-diffusion/model1.safetensors
2024-03-11 10:59:38 INFO [root] Loading: ~/Downloads/stable-diffusion-webui/models/Stable-diffusion/model2.safetensors
2024-03-11 10:59:38 INFO [root] start merging with tensor_sum method
2024-03-11 10:59:47 INFO [root] Init rebasin iterations
2024-03-11 10:59:47 INFO [root] Rebasin iteration 0
stage 1: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 1131/1131 [00:02<00:00, 504.26it/s]
stage 2: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 1133/1133 [00:00<00:00, 177338.75it/s]
2024-03-11 11:00:00 INFO [root] Rebasin iteration 1
stage 1: 100%|██████████████████████████████████████████████████████████████████████████████████████████▊| 1131/1133 [00:00<00:00, 1748.68it/s]
*** API error: POST: http://127.0.0.1:7860/bbwm/merge-models {'error': 'IndexError', 'detail': '', 'body': '', 'errors': 'tuple index out of range'}
Traceback (most recent call last):
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/anyio/streams/memory.py", line 98, in receive
return self.receive_nowait()
^^^^^^^^^^^^^^^^^^^^^
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/anyio/streams/memory.py", line 93, in receive_nowait
raise WouldBlock
anyio.WouldBlock
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/starlette/middleware/base.py", line 78, in call_next
message = await recv_stream.receive()
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/anyio/streams/memory.py", line 118, in receive
raise EndOfStream
anyio.EndOfStream
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "~/Downloads/stable-diffusion-webui/modules/api/api.py", line 186, in exception_handling
return await call_next(request)
^^^^^^^^^^^^^^^^^^^^^^^^
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/starlette/middleware/base.py", line 84, in call_next
raise app_exc
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/starlette/middleware/base.py", line 70, in coro
await self.app(scope, receive_or_disconnect, send_no_error)
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/starlette/middleware/base.py", line 108, in __call__
response = await self.dispatch_func(request, call_next)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "~/Downloads/stable-diffusion-webui/modules/api/api.py", line 150, in log_and_time
res: Response = await call_next(req)
^^^^^^^^^^^^^^^^^^^^
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/starlette/middleware/base.py", line 84, in call_next
raise app_exc
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/starlette/middleware/base.py", line 70, in coro
await self.app(scope, receive_or_disconnect, send_no_error)
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/starlette/middleware/cors.py", line 84, in __call__
await self.app(scope, receive, send)
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/starlette/middleware/gzip.py", line 24, in __call__
await responder(scope, receive, send)
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/starlette/middleware/gzip.py", line 44, in __call__
await self.app(scope, receive, self.send_with_gzip)
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/starlette/middleware/exceptions.py", line 79, in __call__
raise exc
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/starlette/middleware/exceptions.py", line 68, in __call__
await self.app(scope, receive, sender)
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/fastapi/middleware/asyncexitstack.py", line 21, in __call__
raise e
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/fastapi/middleware/asyncexitstack.py", line 18, in __call__
await self.app(scope, receive, send)
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/starlette/routing.py", line 718, in __call__
await route.handle(scope, receive, send)
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/starlette/routing.py", line 276, in handle
await self.app(scope, receive, send)
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/starlette/routing.py", line 66, in app
response = await func(request)
^^^^^^^^^^^^^^^^^^^
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/fastapi/routing.py", line 237, in app
raw_response = await run_endpoint_function(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/fastapi/routing.py", line 163, in run_endpoint_function
return await dependant.call(**values)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "~/Downloads/stable-diffusion-webui/extensions/sd-webui-bayesian-merger/scripts/api.py", line 78, in merge_models_api
merged = merge_models(
^^^^^^^^^^^^^
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/sd_meh/merge.py", line 149, in merge_models
merged = rebasin_merge(
^^^^^^^^^^^^^^
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/sd_meh/merge.py", line 292, in rebasin_merge
thetas["model_a"] = simple_merge(
^^^^^^^^^^^^^
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/sd_meh/merge.py", line 245, in simple_merge
res.result()
File "~/.asdf/installs/python/3.11.5/lib/python3.11/concurrent/futures/_base.py", line 449, in result
return self.__get_result()
^^^^^^^^^^^^^^^^^^^
File "~/.asdf/installs/python/3.11.5/lib/python3.11/concurrent/futures/_base.py", line 401, in __get_result
raise self._exception
File "~/.asdf/installs/python/3.11.5/lib/python3.11/concurrent/futures/thread.py", line 58, in run
result = self.fn(*self.args, **self.kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/sd_meh/merge.py", line 353, in simple_merge_key
with merge_key_context(key, thetas, *args, **kwargs) as result:
File "~/.asdf/installs/python/3.11.5/lib/python3.11/contextlib.py", line 137, in __enter__
return next(self.gen)
^^^^^^^^^^^^^^
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/sd_meh/merge.py", line 450, in merge_key_context
result = merge_key(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/sd_meh/merge.py", line 422, in merge_key
merged_key = merge_method(**merge_args).to(device)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "~/.asdf/installs/python/3.11.5/lib/python3.11/site-packages/sd_meh/merge_methods.py", line 43, in tensor_sum
talphas = int(a.shape[0] * beta)
~~~~~~~^^^
IndexError: tuple index out of range