John6666 commited on
Commit
b242a53
·
verified ·
1 Parent(s): 4a12110

Upload 2 files

Browse files
Files changed (2) hide show
  1. app.py +26 -13
  2. env.py +4 -2
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
- TASK_TYPE_I2I = ["Image-to-Image", "Inpainting", "Flux Fill"] # , "Canny", "Depth"
 
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 == "Canny" or task == "Depth":
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 == "Canny" or task == "Depth": # FluxControlPipeline
123
  if task == "Canny": control_repo = models_canny[0]
124
  elif task == "Depth": control_repo = models_depth[0]
125
- transformer = transformer_model.from_pretrained(control_repo, subfolder="transformer", torch_dtype=dtype, token=hf_token)
126
- text_encoder_2 = t5_model.from_pretrained(control_repo, subfolder="text_encoder_2", torch_dtype=dtype, token=hf_token)
 
 
 
 
 
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 not is_fill: kwargs["strength"] = image_strength
 
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