Upload 2 files

app.py

@@ -18,7 +18,7 @@ import pandas as pd
 import numpy as np
 from pathlib import Path
 
-from env import models, models_dev, models_schnell, models_fill, models_canny, models_depth, num_loras, num_cns, HF_TOKEN, single_file_base_models
+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
 from mod import (clear_cache, get_repo_safetensors, is_repo_name, is_repo_exists, get_model_trigger,
                  description_ui, compose_lora_json, is_valid_lora, fuse_loras, turbo_loras, save_image, preprocess_i2i_image,
                  get_trigger_word, enhance_prompt, set_control_union_image, get_canny_image, get_depth_image,
@@ -62,7 +62,8 @@ last_dtype_str = "BF16"
 
 MAX_SEED = 2**32-1
 TASK_TYPE_T2I = ["Text-to-Image"]
-TASK_TYPE_I2I = ["Image-to-Image", "Inpainting", "Flux Fill"] # , "Canny", "Depth"
+TASK_TYPE_CONTROL = ["Canny", "Depth", "Edit"]
+TASK_TYPE_I2I = ["Image-to-Image", "Inpainting", "Flux Fill"] + TASK_TYPE_CONTROL
 
 def unload_lora():
     global pipe, pipe_i2i
@@ -104,7 +105,7 @@ def load_pipeline(pipe, pipe_i2i, repo_id: str, cn_on: bool, model_type: str, ta
         if task == "Flux Fill":
             pipeline = FluxFillPipeline
             pipeline_i2i = FluxFillPipeline
-        elif task == "Canny" or task == "Depth":
+        elif task in TASK_TYPE_CONTROL:
             pipeline = DiffusionPipeline
             pipeline_i2i = FluxControlPipeline
         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
             print_progress(f"Loading model: {repo_id}", 0, progress)
             pipeline = DiffusionPipeline
             pipeline_i2i = AutoPipelineForInpainting if task == "Inpainting" else AutoPipelineForImage2Image
-        if task == "Canny" or task == "Depth": # FluxControlPipeline
+        if task in TASK_TYPE_CONTROL: # FluxControlPipeline
             if task == "Canny": control_repo = models_canny[0]
             elif task == "Depth": control_repo = models_depth[0]
-            transformer = transformer_model.from_pretrained(control_repo, subfolder="transformer", torch_dtype=dtype, token=hf_token)
-            text_encoder_2 = t5_model.from_pretrained(control_repo, subfolder="text_encoder_2", torch_dtype=dtype, token=hf_token)
+            elif task == "Edit": control_repo = models_edit[0]
+            if task == "Edit":
+                transformer = transformer_model.from_pretrained(control_repo, torch_dtype=dtype, token=hf_token)
+                text_encoder_2 = t5_model.from_pretrained(models_dev[0], subfolder="text_encoder_2", torch_dtype=dtype, token=hf_token)
+            else:
+                transformer = transformer_model.from_pretrained(control_repo, subfolder="transformer", torch_dtype=dtype, token=hf_token)
+                text_encoder_2 = t5_model.from_pretrained(control_repo, subfolder="text_encoder_2", torch_dtype=dtype, token=hf_token)
             #transformer, text_encoder_2 = load_quantized_control(control_repo, dtype, hf_token)
             pipe = pipeline.from_pretrained(models_dev[0], transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=dtype, token=hf_token)
             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
         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)\
             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
         unload_lora()
-        pipe.to("cpu")
-        pipe_i2i.to("cpu")
-        good_vae.to("cpu")
-        taef1.to("cpu")
+        if pipe is not None: pipe.to("cpu")
+        if pipe_i2i is not None: pipe_i2i.to("cpu")
+        if good_vae is not None: good_vae.to("cpu")
+        if taef1 is not None: taef1.to("cpu")
         if controlnet is not None: controlnet.to("cpu")
         if controlnet_union is not None: controlnet_union.to("cpu")
         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
         is_fill = True if task_type == "Flux Fill" else False
         is_depth = True if task_type == "Depth" else False
         is_canny = True if task_type == "Canny" else False
+        is_edit = True if task_type == "Edit" else False
         kwargs = {}
-        if not is_fill: kwargs["strength"] = image_strength
+        if task_type in ["Image-to-Image", "Inpainting"]: kwargs["strength"] = image_strength
+
         if sigmas_factor < 1.0 and task_type != "Image-to-Image": kwargs["sigmas"] = calc_sigmas(steps, sigmas_factor)
 
         with calculateDuration("Generating image"):
@@ -529,6 +537,8 @@ def generate_image_to_image(prompt_mash: str, image_input_path_dict: dict, image
                 pipe_i2i.to(device)
                 pipe_i2i.vae = good_vae
                 image_input = load_image(image_input_path)
+                if task_type in TASK_TYPE_CONTROL: kwargs["control_image"] = image_input
+                else: kwargs["image"] = image_input
                 if is_mask:
                     mask_input = load_image(mask_path)
                     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
                 elif is_depth:
                     image_input = get_depth_image(image_input, height, width)
                     print_progress("Start Depth Inference.")
+                elif is_edit:
+                    print_progress("Start Edit Inference.")
                 else: print_progress("Start I2I Inference.")
                 final_image = pipe_i2i(
                     prompt=prompt_mash,
-                    image=image_input,
+                    #image=image_input,
                     num_inference_steps=steps,
                     guidance_scale=cfg_scale,
                     #width=width,
@@ -559,6 +571,7 @@ def generate_image_to_image(prompt_mash: str, image_input_path_dict: dict, image
                 pipe_i2i.to(device)
                 pipe_i2i.vae = good_vae
                 image_input = load_image(image_input_path)
+                kwargs["image"] = image_input
                 if controlnet_union is not None: controlnet_union.to(device)
                 if controlnet is not None: controlnet.to(device)
                 if is_mask:
@@ -892,12 +905,12 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', fill_width=True, css=css, delete_ca
                                                          layers=False, brush=gr.Brush(colors=["white"], color_mode="fixed", default_size=32), eraser=gr.Eraser(default_size="32"), value=None,
                                                          canvas_size=(384, 384), width=384, height=512)
                         with gr.Column():
-
                             task_type = gr.Radio(label="Task", choices=TASK_TYPE_T2I+TASK_TYPE_I2I, value=TASK_TYPE_T2I[0])
                             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)
                             blur_mask = gr.Checkbox(label="Blur mask", value=False)
                             blur_factor = gr.Slider(label="Blur factor", minimum=0, maximum=50, step=1, value=33)
                             input_image_preprocess = gr.Checkbox(True, label="Preprocess Input image")
+                            gr.Markdown("About ['Edit' Control weights](https://huggingface.co/sayakpaul/edit-control-lr_1e-4-wd_1e-4-gs_15.0-cd_0.1)")
                 with gr.Tab("More LoRA"):
                     with gr.Accordion("External LoRA", open=True):
                         with gr.Column():

env.py

@@ -77,9 +77,11 @@ models_fill = ["fuliucansheng/FLUX.1-Fill-dev-diffusers"]
 
 models_dedistill = []
 
-models_canny = ["sayakpaul/FLUX.1-Canny-dev-nf4"]
+models_canny = ["fuliucansheng/FLUX.1-Canny-dev-diffusers", "sayakpaul/FLUX.1-Canny-dev-nf4"]
 
-models_depth = ["sayakpaul/FLUX.1-Depth-dev-nf4"]
+models_depth = ["fuliucansheng/FLUX.1-Depth-dev-diffusers", "sayakpaul/FLUX.1-Depth-dev-nf4"]
+
+models_edit = ["sayakpaul/edit-control-lr_1e-4-wd_1e-4-gs_15.0-cd_0.1"]
 
 models = models_dev + models_schnell + models_fill