From 0b1ddffe407a0dbf222d016d79ab9fbb2fd18bbf Mon Sep 17 00:00:00 2001
From: Anatoly Belikov <abelikov@singularitynet.io>
Date: Fri, 6 Sep 2024 18:25:12 +0300
Subject: [PATCH] add flux symlinks

---
 .github/workflows/python-app.yml     | 2 +-
 .gitmodules                          | 2 +-
 multigen/loader.py                   | 9 +++++----
 tests/models-sd/flux/tiny-flux-pipe1 | 1 +
 tests/models-sd/flux/tiny-flux-pipe2 | 1 +
 tests/models-sd/flux/tiny-flux-pipe3 | 1 +
 tests/models-sd/flux/tiny-flux-pipe4 | 1 +
 7 files changed, 11 insertions(+), 6 deletions(-)
 create mode 120000 tests/models-sd/flux/tiny-flux-pipe1
 create mode 120000 tests/models-sd/flux/tiny-flux-pipe2
 create mode 120000 tests/models-sd/flux/tiny-flux-pipe3
 create mode 120000 tests/models-sd/flux/tiny-flux-pipe4

diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml
index 2513f2b..3d4c0a1 100644
--- a/.github/workflows/python-app.yml
+++ b/.github/workflows/python-app.yml
@@ -34,7 +34,7 @@ jobs:
         cd tests/models-sd/SDXL/tiny-sdxl && git lfs pull
     - name: pull lfs
       run: |
-        cd tests/models-sd/tiny-flux-pipe && git lfs pull
+        cd tests/models-sd/flux/tiny-flux-pipe && git lfs pull
     - name: symlink models directory
       run: ln -s ../../../../models tests/models-full
     - name: Test with pytest
diff --git a/.gitmodules b/.gitmodules
index 46ef9a2..8e7f44c 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -4,6 +4,6 @@
 [submodule "tests/models-sd/SDXL/tiny-sdxl"]
 	path = tests/models-sd/SDXL/tiny-sdxl
 	url = https://huggingface.co/hf-internal-testing/tiny-sdxl-pipe
-[submodule "tests/models-sd/tiny-flux-pipe"]
+[submodule "tests/models-sd/flux/tiny-flux-pipe"]
 	path = tests/models-sd/flux/tiny-flux-pipe
 	url = https://huggingface.co/hf-internal-testing/tiny-flux-pipe
diff --git a/multigen/loader.py b/multigen/loader.py
index 6430523..8dd8e97 100644
--- a/multigen/loader.py
+++ b/multigen/loader.py
@@ -1,6 +1,6 @@
 from typing import Type, List
 import random
-import copy
+import copy as cp
 from contextlib import nullcontext
 import torch
 import logging
@@ -28,10 +28,11 @@ def weightshare_copy(pipe):
     ctx = init_empty_weights if is_accelerate_available() else nullcontext
     with ctx():
         for key, component in copy.components.items():
-            if key in ('tokenizer', 'tokenizer_2', 'feature_extractor'):
-                continue
             if getattr(copy, key) is None:
                 continue
+            if key in ('tokenizer', 'tokenizer_2', 'feature_extractor'):
+                setattr(copy, key, cp.deepcopy(getattr(copy, key)))
+                continue
             cls = getattr(copy, key).__class__
             if hasattr(cls, 'from_config'):
                 setattr(copy, key, cls.from_config(getattr(copy, key).config))
@@ -141,7 +142,7 @@ def cache_pipeline(self, pipe: DiffusionPipeline, model_id):
                     self._cpu_pipes.pop(key_to_delete)
                 item = pipe
                 if pipe.device.type == 'cuda':
-                    item = copy.deepcopy(pipe).to('cpu')
+                    item = cp.deepcopy(pipe).to('cpu')
                 self._cpu_pipes[model_id] = item
                 logger.debug(f'storing {model_id} on cpu')
             assert pipe.device == device
diff --git a/tests/models-sd/flux/tiny-flux-pipe1 b/tests/models-sd/flux/tiny-flux-pipe1
new file mode 120000
index 0000000..146f84b
--- /dev/null
+++ b/tests/models-sd/flux/tiny-flux-pipe1
@@ -0,0 +1 @@
+tiny-flux-pipe/
\ No newline at end of file
diff --git a/tests/models-sd/flux/tiny-flux-pipe2 b/tests/models-sd/flux/tiny-flux-pipe2
new file mode 120000
index 0000000..146f84b
--- /dev/null
+++ b/tests/models-sd/flux/tiny-flux-pipe2
@@ -0,0 +1 @@
+tiny-flux-pipe/
\ No newline at end of file
diff --git a/tests/models-sd/flux/tiny-flux-pipe3 b/tests/models-sd/flux/tiny-flux-pipe3
new file mode 120000
index 0000000..146f84b
--- /dev/null
+++ b/tests/models-sd/flux/tiny-flux-pipe3
@@ -0,0 +1 @@
+tiny-flux-pipe/
\ No newline at end of file
diff --git a/tests/models-sd/flux/tiny-flux-pipe4 b/tests/models-sd/flux/tiny-flux-pipe4
new file mode 120000
index 0000000..146f84b
--- /dev/null
+++ b/tests/models-sd/flux/tiny-flux-pipe4
@@ -0,0 +1 @@
+tiny-flux-pipe/
\ No newline at end of file