File size: 12,756 Bytes
9ceef7d
09fa6ac
 
9ceef7d
cd39c08
09fa6ac
 
f3a071e
09fa6ac
 
c60bd9d
09fa6ac
 
 
 
 
f1d6334
 
 
cd39c08
dc1b7f5
cd39c08
 
39167cc
cd39c08
 
 
 
 
 
 
 
 
5e404f6
c9be37f
cd39c08
 
 
4c19db8
 
 
 
 
 
 
 
 
 
 
09fa6ac
f0ac7fb
 
 
245e508
f0ac7fb
 
 
5e404f6
09fa6ac
 
f0c3651
0518209
f0c3651
 
 
 
 
 
 
 
 
 
 
 
 
09fa6ac
 
 
 
cd39c08
09fa6ac
 
c60bd9d
09fa6ac
5e404f6
09fa6ac
 
 
 
 
 
dc1b7f5
 
 
 
 
 
 
 
 
 
 
 
 
 
f1d6334
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8ac6201
f1d6334
 
 
 
 
 
 
 
 
 
 
 
 
 
8ac6201
f1d6334
8ac6201
f1d6334
 
 
 
 
 
8ac6201
f1d6334
 
 
 
 
5302530
467b9b3
dc1b7f5
f1d6334
a138792
f1d6334
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dc1b7f5
f1d6334
 
 
 
 
dc1b7f5
f1d6334
dc1b7f5
d0d2198
 
f1d6334
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d0d2198
f1d6334
d0d2198
f1d6334
 
 
 
cd39c08
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f1d6334
f3a071e
 
 
 
 
 
eb7cad2
cd39c08
467b9b3
5e404f6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cd39c08
 
 
 
 
 
8ac6201
f1d6334
cd39c08
 
c60bd9d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1553d93
c60bd9d
4c19db8
c60bd9d
 
 
 
f3a071e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c60bd9d
8ac6201
f0c3651
 
f0ac7fb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
import spaces
import gradio as gr
import torch
from PIL import Image
from pathlib import Path
import gc
import subprocess
from env import num_cns, model_trigger


subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
subprocess.run('pip cache purge', shell=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
torch.set_grad_enabled(False)


control_images = [None] * num_cns
control_modes = [-1] * num_cns
control_scales = [0] * num_cns


def is_repo_name(s):
    import re
    return re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', s)


def is_repo_exists(repo_id):
    from huggingface_hub import HfApi
    api = HfApi()
    try:
        if api.repo_exists(repo_id=repo_id): return True
        else: return False
    except Exception as e:
        print(f"Error: Failed to connect {repo_id}.")
        print(e)
        return True # for safe


from translatepy import Translator
translator = Translator()
def translate_to_en(input: str):
    try:
        output = str(translator.translate(input, 'English'))
    except Exception as e:
        output = input
        print(e)
    return output


def clear_cache():
    try:
        torch.cuda.empty_cache()
        #torch.cuda.reset_max_memory_allocated()
        #torch.cuda.reset_peak_memory_stats()
        gc.collect()
    except Exception as e:
        print(e)
        raise Exception(f"Cache clearing error: {e}") from e


def deselect_lora():
    selected_index = None
    new_placeholder = "Type a prompt"
    updated_text = ""
    width = 1024
    height = 1024
    return (
        gr.update(placeholder=new_placeholder),
        updated_text,
        selected_index,
        width,
        height,
    )


def get_repo_safetensors(repo_id: str):
    from huggingface_hub import HfApi
    api = HfApi()
    try:
        if not is_repo_name(repo_id) or not is_repo_exists(repo_id): return gr.update(value="", choices=[])
        files = api.list_repo_files(repo_id=repo_id)
    except Exception as e:
        print(f"Error: Failed to get {repo_id}'s info.")
        print(e)
        gr.Warning(f"Error: Failed to get {repo_id}'s info.")
        return gr.update(choices=[])
    files = [f for f in files if f.endswith(".safetensors")]
    if len(files) == 0: return gr.update(value="", choices=[])
    else: return gr.update(value=files[0], choices=files)


def expand2square(pil_img: Image.Image, background_color: tuple=(0, 0, 0)):
    width, height = pil_img.size
    if width == height:
        return pil_img
    elif width > height:
        result = Image.new(pil_img.mode, (width, width), background_color)
        result.paste(pil_img, (0, (width - height) // 2))
        return result
    else:
        result = Image.new(pil_img.mode, (height, height), background_color)
        result.paste(pil_img, ((height - width) // 2, 0))
        return result


# https://huggingface.co/spaces/DamarJati/FLUX.1-DEV-Canny/blob/main/app.py
def resize_image(image, target_width, target_height, crop=True):
    from image_datasets.canny_dataset import c_crop
    if crop:
        image = c_crop(image)  # Crop the image to square
        original_width, original_height = image.size

        # Resize to match the target size without stretching
        scale = max(target_width / original_width, target_height / original_height)
        resized_width = int(scale * original_width)
        resized_height = int(scale * original_height)

        image = image.resize((resized_width, resized_height), Image.LANCZOS)
        
        # Center crop to match the target dimensions
        left = (resized_width - target_width) // 2
        top = (resized_height - target_height) // 2
        image = image.crop((left, top, left + target_width, top + target_height))
    else:
        image = image.resize((target_width, target_height), Image.LANCZOS)
    
    return image


# https://huggingface.co/spaces/jiuface/FLUX.1-dev-Controlnet-Union/blob/main/app.py
# https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union
controlnet_union_modes = {
    "None": -1,
    #"scribble_hed": 0,
    "canny": 0, # supported
    "mlsd": 0, #supported
    "tile": 1, #supported
    "depth_midas": 2, # supported
    "blur": 3, # supported
    "openpose": 4,  # supported
    "gray": 5,  # supported
    "low_quality": 6,  # supported
}


# https://github.com/pytorch/pytorch/issues/123834
def get_control_params():
    from diffusers.utils import load_image
    modes = []
    images = []
    scales = []
    for i, mode in enumerate(control_modes):
        if mode == -1 or control_images[i] is None: continue
        modes.append(control_modes[i])
        images.append(load_image(control_images[i]))
        scales.append(control_scales[i])
    return modes, images, scales


from preprocessor import Preprocessor
def preprocess_image(image: Image.Image, control_mode: str, height: int, width: int,

                     preprocess_resolution: int):
    if control_mode == "None": return image
    image_resolution = max(width, height)
    image_before = resize_image(expand2square(image.convert("RGB")), image_resolution, image_resolution, False)
    # generated control_
    print("start to generate control image")
    preprocessor = Preprocessor()
    if control_mode == "depth_midas":
        preprocessor.load("Midas")
        control_image = preprocessor(
            image=image_before,
            image_resolution=image_resolution,
            detect_resolution=preprocess_resolution,
        )
    if control_mode == "openpose":
        preprocessor.load("Openpose")
        control_image = preprocessor(
            image=image_before,
            hand_and_face=True,
            image_resolution=image_resolution,
            detect_resolution=preprocess_resolution,
        )
    if control_mode == "canny":
        preprocessor.load("Canny")
        control_image = preprocessor(
            image=image_before,
            image_resolution=image_resolution,
            detect_resolution=preprocess_resolution,
        )

    if control_mode == "mlsd":
        preprocessor.load("MLSD")
        control_image = preprocessor(
            image=image_before,
            image_resolution=image_resolution,
            detect_resolution=preprocess_resolution,
        )

    if control_mode == "scribble_hed":
        preprocessor.load("HED")
        control_image = preprocessor(
            image=image_before,
            image_resolution=image_resolution,
            detect_resolution=preprocess_resolution,
        )
    
    if control_mode == "low_quality" or control_mode == "gray" or control_mode == "blur" or control_mode == "tile":
        control_image = image_before
        image_width = 768
        image_height = 768
    else:
        # make sure control image size is same as resized_image
        image_width, image_height = control_image.size
    
    image_after = resize_image(control_image, width, height, False)
    ref_width, ref_height = image.size
    print(f"generate control image success: {ref_width}x{ref_height} => {image_width}x{image_height}")
    return image_after


def get_control_union_mode():
    return list(controlnet_union_modes.keys())


def set_control_union_mode(i: int, mode: str, scale: str):
    global control_modes
    global control_scales
    control_modes[i] = controlnet_union_modes.get(mode, 0)
    control_scales[i] = scale
    if mode != "None": return True
    else: return gr.update(visible=True)


def set_control_union_image(i: int, mode: str, image: Image.Image | None, height: int, width: int, preprocess_resolution: int):
    global control_images
    if image is None: return None
    control_images[i] = preprocess_image(image, mode, height, width, preprocess_resolution)
    return control_images[i]


def compose_lora_json(lorajson: list[dict], i: int, name: str, scale: float, filename: str, trigger: str):
    lorajson[i]["name"] = str(name) if name != "None" else ""
    lorajson[i]["scale"] = float(scale)
    lorajson[i]["filename"] = str(filename)
    lorajson[i]["trigger"] = str(trigger)
    return lorajson


def is_valid_lora(lorajson: list[dict]):
    valid = False
    for d in lorajson:
        if "name" in d.keys() and d["name"] and d["name"] != "None": valid = True
    return valid


def get_trigger_word(lorajson: list[dict]):
    trigger = ""
    for d in lorajson:
        if "name" in d.keys() and d["name"] and d["name"] != "None" and d["trigger"]:
            trigger += ", " + d["trigger"]
    return trigger


def get_model_trigger(model_name: str):
    trigger = ""
    if model_name in model_trigger.keys(): trigger += ", " + model_trigger[model_name]
    return trigger


# https://huggingface.co/docs/diffusers/v0.23.1/en/api/loaders#diffusers.loaders.LoraLoaderMixin.fuse_lora
# https://github.com/huggingface/diffusers/issues/4919
def fuse_loras(pipe, lorajson: list[dict]):
    try:
        if not lorajson or not isinstance(lorajson, list): return
        a_list = []
        w_list = []
        for d in lorajson:
            if not d or not isinstance(d, dict) or not d["name"] or d["name"] == "None": continue
            k = d["name"]
            if is_repo_name(k) and is_repo_exists(k):
                a_name = Path(k).stem
                pipe.load_lora_weights(k, weight_name=d["filename"], adapter_name = a_name)
            elif not Path(k).exists():
                print(f"LoRA not found: {k}")
                continue
            else:
                w_name = Path(k).name
                a_name = Path(k).stem
                pipe.load_lora_weights(k, weight_name = w_name, adapter_name = a_name)
            a_list.append(a_name)
            w_list.append(d["scale"])
        if not a_list: return
        pipe.set_adapters(a_list, adapter_weights=w_list)
        pipe.fuse_lora(adapter_names=a_list, lora_scale=1.0)
        #pipe.unload_lora_weights()
    except Exception as e:
        print(f"External LoRA Error: {e}")
        raise Exception(f"External LoRA Error: {e}") from e


def description_ui():
    gr.Markdown(
        """

- Mod of [multimodalart/flux-lora-the-explorer](https://huggingface.co/spaces/multimodalart/flux-lora-the-explorer),

 [jiuface/FLUX.1-dev-Controlnet-Union](https://huggingface.co/spaces/jiuface/FLUX.1-dev-Controlnet-Union),

 [DamarJati/FLUX.1-DEV-Canny](https://huggingface.co/spaces/DamarJati/FLUX.1-DEV-Canny),

 [gokaygokay/FLUX-Prompt-Generator](https://huggingface.co/spaces/gokaygokay/FLUX-Prompt-Generator).

"""
    )


from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
def load_prompt_enhancer():
    try:
        model_checkpoint = "gokaygokay/Flux-Prompt-Enhance"
        tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint).eval().to(device=device)
        enhancer_flux = pipeline('text2text-generation', model=model, tokenizer=tokenizer, repetition_penalty=1.5, device=device)
    except Exception as e:
        print(e)
        enhancer_flux = None
    return enhancer_flux


enhancer_flux = load_prompt_enhancer()


@spaces.GPU(duration=30)
def enhance_prompt(input_prompt):
    result = enhancer_flux("enhance prompt: " + translate_to_en(input_prompt), max_length = 256)
    enhanced_text = result[0]['generated_text']
    return enhanced_text


def save_image(image, savefile, modelname, prompt, height, width, steps, cfg, seed):
    import uuid
    from PIL import Image, PngImagePlugin
    import json
    try:
        if savefile is None: savefile = f"{modelname.split('/')[-1]}_{str(uuid.uuid4())}.png"
        metadata = {"prompt": prompt, "Model": {"Model": modelname.split("/")[-1]}}
        metadata["num_inference_steps"] = steps
        metadata["guidance_scale"] = cfg
        metadata["seed"] = seed
        metadata["resolution"] = f"{width} x {height}"
        metadata_str = json.dumps(metadata)
        info = PngImagePlugin.PngInfo()
        info.add_text("metadata", metadata_str)
        image.save(savefile, "PNG", pnginfo=info)
        return str(Path(savefile).resolve())
    except Exception as e:
        print(f"Failed to save image file: {e}")
        raise Exception(f"Failed to save image file:") from e


load_prompt_enhancer.zerogpu = True
fuse_loras.zerogpu = True
preprocess_image.zerogpu = True
get_control_params.zerogpu = True
clear_cache.zerogpu = True