Spaces:
Running
on
Zero
Running
on
Zero
Upload 2 files
Browse files
app.py
CHANGED
@@ -18,7 +18,7 @@ import pandas as pd
|
|
18 |
import numpy as np
|
19 |
from pathlib import Path
|
20 |
|
21 |
-
from env import models, models_dev, models_schnell, models_fill, models_canny, models_depth, num_loras, num_cns, HF_TOKEN, single_file_base_models
|
22 |
from mod import (clear_cache, get_repo_safetensors, is_repo_name, is_repo_exists, get_model_trigger,
|
23 |
description_ui, compose_lora_json, is_valid_lora, fuse_loras, turbo_loras, save_image, preprocess_i2i_image,
|
24 |
get_trigger_word, enhance_prompt, set_control_union_image, get_canny_image, get_depth_image,
|
@@ -62,7 +62,8 @@ last_dtype_str = "BF16"
|
|
62 |
|
63 |
MAX_SEED = 2**32-1
|
64 |
TASK_TYPE_T2I = ["Text-to-Image"]
|
65 |
-
|
|
|
66 |
|
67 |
def unload_lora():
|
68 |
global pipe, pipe_i2i
|
@@ -104,7 +105,7 @@ def load_pipeline(pipe, pipe_i2i, repo_id: str, cn_on: bool, model_type: str, ta
|
|
104 |
if task == "Flux Fill":
|
105 |
pipeline = FluxFillPipeline
|
106 |
pipeline_i2i = FluxFillPipeline
|
107 |
-
elif task
|
108 |
pipeline = DiffusionPipeline
|
109 |
pipeline_i2i = FluxControlPipeline
|
110 |
elif cn_on: # with ControlNet
|
@@ -119,11 +120,16 @@ def load_pipeline(pipe, pipe_i2i, repo_id: str, cn_on: bool, model_type: str, ta
|
|
119 |
print_progress(f"Loading model: {repo_id}", 0, progress)
|
120 |
pipeline = DiffusionPipeline
|
121 |
pipeline_i2i = AutoPipelineForInpainting if task == "Inpainting" else AutoPipelineForImage2Image
|
122 |
-
if task
|
123 |
if task == "Canny": control_repo = models_canny[0]
|
124 |
elif task == "Depth": control_repo = models_depth[0]
|
125 |
-
|
126 |
-
|
|
|
|
|
|
|
|
|
|
|
127 |
#transformer, text_encoder_2 = load_quantized_control(control_repo, dtype, hf_token)
|
128 |
pipe = pipeline.from_pretrained(models_dev[0], transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=dtype, token=hf_token)
|
129 |
pipe_i2i = pipeline_i2i.from_pipe(pipe, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=dtype)
|
@@ -161,10 +167,10 @@ def change_base_model(repo_id: str, cn_on: bool, disable_model_cache: bool, mode
|
|
161 |
if not disable_model_cache and (repo_id == last_model and cn_on is last_cn_on and task == last_task and dtype_str == last_dtype_str)\
|
162 |
or ((not is_repo_name(repo_id) or not is_repo_exists(repo_id)) and not ".safetensors" in repo_id): return gr.update() # and not ".gguf" in repo_id
|
163 |
unload_lora()
|
164 |
-
pipe.to("cpu")
|
165 |
-
pipe_i2i.to("cpu")
|
166 |
-
good_vae.to("cpu")
|
167 |
-
taef1.to("cpu")
|
168 |
if controlnet is not None: controlnet.to("cpu")
|
169 |
if controlnet_union is not None: controlnet_union.to("cpu")
|
170 |
pipe, pipe_i2i = load_pipeline(pipe, pipe_i2i, repo_id, cn_on, model_type, task, dtype_str, HF_TOKEN, progress)
|
@@ -518,8 +524,10 @@ def generate_image_to_image(prompt_mash: str, image_input_path_dict: dict, image
|
|
518 |
is_fill = True if task_type == "Flux Fill" else False
|
519 |
is_depth = True if task_type == "Depth" else False
|
520 |
is_canny = True if task_type == "Canny" else False
|
|
|
521 |
kwargs = {}
|
522 |
-
if
|
|
|
523 |
if sigmas_factor < 1.0 and task_type != "Image-to-Image": kwargs["sigmas"] = calc_sigmas(steps, sigmas_factor)
|
524 |
|
525 |
with calculateDuration("Generating image"):
|
@@ -529,6 +537,8 @@ def generate_image_to_image(prompt_mash: str, image_input_path_dict: dict, image
|
|
529 |
pipe_i2i.to(device)
|
530 |
pipe_i2i.vae = good_vae
|
531 |
image_input = load_image(image_input_path)
|
|
|
|
|
532 |
if is_mask:
|
533 |
mask_input = load_image(mask_path)
|
534 |
if blur_mask: mask_input = pipe_i2i.mask_processor.blur(mask_input, blur_factor=blur_factor)
|
@@ -541,10 +551,12 @@ def generate_image_to_image(prompt_mash: str, image_input_path_dict: dict, image
|
|
541 |
elif is_depth:
|
542 |
image_input = get_depth_image(image_input, height, width)
|
543 |
print_progress("Start Depth Inference.")
|
|
|
|
|
544 |
else: print_progress("Start I2I Inference.")
|
545 |
final_image = pipe_i2i(
|
546 |
prompt=prompt_mash,
|
547 |
-
image=image_input,
|
548 |
num_inference_steps=steps,
|
549 |
guidance_scale=cfg_scale,
|
550 |
#width=width,
|
@@ -559,6 +571,7 @@ def generate_image_to_image(prompt_mash: str, image_input_path_dict: dict, image
|
|
559 |
pipe_i2i.to(device)
|
560 |
pipe_i2i.vae = good_vae
|
561 |
image_input = load_image(image_input_path)
|
|
|
562 |
if controlnet_union is not None: controlnet_union.to(device)
|
563 |
if controlnet is not None: controlnet.to(device)
|
564 |
if is_mask:
|
@@ -892,12 +905,12 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', fill_width=True, css=css, delete_ca
|
|
892 |
layers=False, brush=gr.Brush(colors=["white"], color_mode="fixed", default_size=32), eraser=gr.Eraser(default_size="32"), value=None,
|
893 |
canvas_size=(384, 384), width=384, height=512)
|
894 |
with gr.Column():
|
895 |
-
|
896 |
task_type = gr.Radio(label="Task", choices=TASK_TYPE_T2I+TASK_TYPE_I2I, value=TASK_TYPE_T2I[0])
|
897 |
image_strength = gr.Slider(label="Strength", info="Lower means more image influence in I2I, opposite in Inpaint", minimum=0.01, maximum=1.0, step=0.01, value=0.75)
|
898 |
blur_mask = gr.Checkbox(label="Blur mask", value=False)
|
899 |
blur_factor = gr.Slider(label="Blur factor", minimum=0, maximum=50, step=1, value=33)
|
900 |
input_image_preprocess = gr.Checkbox(True, label="Preprocess Input image")
|
|
|
901 |
with gr.Tab("More LoRA"):
|
902 |
with gr.Accordion("External LoRA", open=True):
|
903 |
with gr.Column():
|
|
|
18 |
import numpy as np
|
19 |
from pathlib import Path
|
20 |
|
21 |
+
from env import models, models_dev, models_schnell, models_fill, models_canny, models_depth, models_edit, num_loras, num_cns, HF_TOKEN, single_file_base_models
|
22 |
from mod import (clear_cache, get_repo_safetensors, is_repo_name, is_repo_exists, get_model_trigger,
|
23 |
description_ui, compose_lora_json, is_valid_lora, fuse_loras, turbo_loras, save_image, preprocess_i2i_image,
|
24 |
get_trigger_word, enhance_prompt, set_control_union_image, get_canny_image, get_depth_image,
|
|
|
62 |
|
63 |
MAX_SEED = 2**32-1
|
64 |
TASK_TYPE_T2I = ["Text-to-Image"]
|
65 |
+
TASK_TYPE_CONTROL = ["Canny", "Depth", "Edit"]
|
66 |
+
TASK_TYPE_I2I = ["Image-to-Image", "Inpainting", "Flux Fill"] + TASK_TYPE_CONTROL
|
67 |
|
68 |
def unload_lora():
|
69 |
global pipe, pipe_i2i
|
|
|
105 |
if task == "Flux Fill":
|
106 |
pipeline = FluxFillPipeline
|
107 |
pipeline_i2i = FluxFillPipeline
|
108 |
+
elif task in TASK_TYPE_CONTROL:
|
109 |
pipeline = DiffusionPipeline
|
110 |
pipeline_i2i = FluxControlPipeline
|
111 |
elif cn_on: # with ControlNet
|
|
|
120 |
print_progress(f"Loading model: {repo_id}", 0, progress)
|
121 |
pipeline = DiffusionPipeline
|
122 |
pipeline_i2i = AutoPipelineForInpainting if task == "Inpainting" else AutoPipelineForImage2Image
|
123 |
+
if task in TASK_TYPE_CONTROL: # FluxControlPipeline
|
124 |
if task == "Canny": control_repo = models_canny[0]
|
125 |
elif task == "Depth": control_repo = models_depth[0]
|
126 |
+
elif task == "Edit": control_repo = models_edit[0]
|
127 |
+
if task == "Edit":
|
128 |
+
transformer = transformer_model.from_pretrained(control_repo, torch_dtype=dtype, token=hf_token)
|
129 |
+
text_encoder_2 = t5_model.from_pretrained(models_dev[0], subfolder="text_encoder_2", torch_dtype=dtype, token=hf_token)
|
130 |
+
else:
|
131 |
+
transformer = transformer_model.from_pretrained(control_repo, subfolder="transformer", torch_dtype=dtype, token=hf_token)
|
132 |
+
text_encoder_2 = t5_model.from_pretrained(control_repo, subfolder="text_encoder_2", torch_dtype=dtype, token=hf_token)
|
133 |
#transformer, text_encoder_2 = load_quantized_control(control_repo, dtype, hf_token)
|
134 |
pipe = pipeline.from_pretrained(models_dev[0], transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=dtype, token=hf_token)
|
135 |
pipe_i2i = pipeline_i2i.from_pipe(pipe, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=dtype)
|
|
|
167 |
if not disable_model_cache and (repo_id == last_model and cn_on is last_cn_on and task == last_task and dtype_str == last_dtype_str)\
|
168 |
or ((not is_repo_name(repo_id) or not is_repo_exists(repo_id)) and not ".safetensors" in repo_id): return gr.update() # and not ".gguf" in repo_id
|
169 |
unload_lora()
|
170 |
+
if pipe is not None: pipe.to("cpu")
|
171 |
+
if pipe_i2i is not None: pipe_i2i.to("cpu")
|
172 |
+
if good_vae is not None: good_vae.to("cpu")
|
173 |
+
if taef1 is not None: taef1.to("cpu")
|
174 |
if controlnet is not None: controlnet.to("cpu")
|
175 |
if controlnet_union is not None: controlnet_union.to("cpu")
|
176 |
pipe, pipe_i2i = load_pipeline(pipe, pipe_i2i, repo_id, cn_on, model_type, task, dtype_str, HF_TOKEN, progress)
|
|
|
524 |
is_fill = True if task_type == "Flux Fill" else False
|
525 |
is_depth = True if task_type == "Depth" else False
|
526 |
is_canny = True if task_type == "Canny" else False
|
527 |
+
is_edit = True if task_type == "Edit" else False
|
528 |
kwargs = {}
|
529 |
+
if task_type in ["Image-to-Image", "Inpainting"]: kwargs["strength"] = image_strength
|
530 |
+
|
531 |
if sigmas_factor < 1.0 and task_type != "Image-to-Image": kwargs["sigmas"] = calc_sigmas(steps, sigmas_factor)
|
532 |
|
533 |
with calculateDuration("Generating image"):
|
|
|
537 |
pipe_i2i.to(device)
|
538 |
pipe_i2i.vae = good_vae
|
539 |
image_input = load_image(image_input_path)
|
540 |
+
if task_type in TASK_TYPE_CONTROL: kwargs["control_image"] = image_input
|
541 |
+
else: kwargs["image"] = image_input
|
542 |
if is_mask:
|
543 |
mask_input = load_image(mask_path)
|
544 |
if blur_mask: mask_input = pipe_i2i.mask_processor.blur(mask_input, blur_factor=blur_factor)
|
|
|
551 |
elif is_depth:
|
552 |
image_input = get_depth_image(image_input, height, width)
|
553 |
print_progress("Start Depth Inference.")
|
554 |
+
elif is_edit:
|
555 |
+
print_progress("Start Edit Inference.")
|
556 |
else: print_progress("Start I2I Inference.")
|
557 |
final_image = pipe_i2i(
|
558 |
prompt=prompt_mash,
|
559 |
+
#image=image_input,
|
560 |
num_inference_steps=steps,
|
561 |
guidance_scale=cfg_scale,
|
562 |
#width=width,
|
|
|
571 |
pipe_i2i.to(device)
|
572 |
pipe_i2i.vae = good_vae
|
573 |
image_input = load_image(image_input_path)
|
574 |
+
kwargs["image"] = image_input
|
575 |
if controlnet_union is not None: controlnet_union.to(device)
|
576 |
if controlnet is not None: controlnet.to(device)
|
577 |
if is_mask:
|
|
|
905 |
layers=False, brush=gr.Brush(colors=["white"], color_mode="fixed", default_size=32), eraser=gr.Eraser(default_size="32"), value=None,
|
906 |
canvas_size=(384, 384), width=384, height=512)
|
907 |
with gr.Column():
|
|
|
908 |
task_type = gr.Radio(label="Task", choices=TASK_TYPE_T2I+TASK_TYPE_I2I, value=TASK_TYPE_T2I[0])
|
909 |
image_strength = gr.Slider(label="Strength", info="Lower means more image influence in I2I, opposite in Inpaint", minimum=0.01, maximum=1.0, step=0.01, value=0.75)
|
910 |
blur_mask = gr.Checkbox(label="Blur mask", value=False)
|
911 |
blur_factor = gr.Slider(label="Blur factor", minimum=0, maximum=50, step=1, value=33)
|
912 |
input_image_preprocess = gr.Checkbox(True, label="Preprocess Input image")
|
913 |
+
gr.Markdown("About ['Edit' Control weights](https://huggingface.co/sayakpaul/edit-control-lr_1e-4-wd_1e-4-gs_15.0-cd_0.1)")
|
914 |
with gr.Tab("More LoRA"):
|
915 |
with gr.Accordion("External LoRA", open=True):
|
916 |
with gr.Column():
|
env.py
CHANGED
@@ -77,9 +77,11 @@ models_fill = ["fuliucansheng/FLUX.1-Fill-dev-diffusers"]
|
|
77 |
|
78 |
models_dedistill = []
|
79 |
|
80 |
-
models_canny = ["sayakpaul/FLUX.1-Canny-dev-nf4"]
|
81 |
|
82 |
-
models_depth = ["sayakpaul/FLUX.1-Depth-dev-nf4"]
|
|
|
|
|
83 |
|
84 |
models = models_dev + models_schnell + models_fill
|
85 |
|
|
|
77 |
|
78 |
models_dedistill = []
|
79 |
|
80 |
+
models_canny = ["fuliucansheng/FLUX.1-Canny-dev-diffusers", "sayakpaul/FLUX.1-Canny-dev-nf4"]
|
81 |
|
82 |
+
models_depth = ["fuliucansheng/FLUX.1-Depth-dev-diffusers", "sayakpaul/FLUX.1-Depth-dev-nf4"]
|
83 |
+
|
84 |
+
models_edit = ["sayakpaul/edit-control-lr_1e-4-wd_1e-4-gs_15.0-cd_0.1"]
|
85 |
|
86 |
models = models_dev + models_schnell + models_fill
|
87 |
|