Update app.py

app.py

@@ -11,7 +11,11 @@ from diffusers import (
     DDIMScheduler,
     UNet2DConditionModel,
 )
-from modules.model import CrossAttnProcessor, StableDiffusionPipeline, load_lora_attn_procs
+from modules.model_pww import (
+    CrossAttnProcessor,
+    StableDiffusionPipeline,
+    load_lora_attn_procs,
+)
 from torchvision import transforms
 from transformers import CLIPTokenizer, CLIPTextModel
 from PIL import Image
@@ -20,11 +24,15 @@ from safetensors.torch import load_file
 import modules.safe as _
 
 models = [
-    ("AbyssOrangeMix_Base", "OrangeMix/AbyssOrangeMix2"),
+    ("AbyssOrangeMix2", "Korakoe/AbyssOrangeMix2-HF"),
+    ("Anything 4.0", "andite/anything-v4.0"),
+    ("Open Journey", "prompthero/openjourney"),
+    ("Basil Mix", "nuigurumi/basil_mix"),
+    ("ACertainModel", "JosephusCheung/ACertainModel"),
 ]
 
-base_name = "AbyssOrangeMix_Base"
-base_model = "OrangeMix/AbyssOrangeMix2"
+base_name, base_model = models[0]
+clip_skip = 2
 
 samplers_k_diffusion = [
     ("Euler a", "sample_euler_ancestral", {}),
@@ -36,24 +44,20 @@ samplers_k_diffusion = [
     ("DPM++ 2S a", "sample_dpmpp_2s_ancestral", {}),
     ("DPM++ 2M", "sample_dpmpp_2m", {}),
     ("DPM++ SDE", "sample_dpmpp_sde", {}),
-    ("DPM fast", "sample_dpm_fast", {}),
-    ("DPM adaptive", "sample_dpm_adaptive", {}),
     ("LMS Karras", "sample_lms", {"scheduler": "karras"}),
-    (
-        "DPM2 Karras",
-        "sample_dpm_2",
-        {"scheduler": "karras", "discard_next_to_last_sigma": True},
-    ),
-    (
-        "DPM2 a Karras",
-        "sample_dpm_2_ancestral",
-        {"scheduler": "karras", "discard_next_to_last_sigma": True},
-    ),
+    ("DPM2 Karras", "sample_dpm_2", {"scheduler": "karras", "discard_next_to_last_sigma": True}),
+    ("DPM2 a Karras", "sample_dpm_2_ancestral", {"scheduler": "karras", "discard_next_to_last_sigma": True}),
     ("DPM++ 2S a Karras", "sample_dpmpp_2s_ancestral", {"scheduler": "karras"}),
     ("DPM++ 2M Karras", "sample_dpmpp_2m", {"scheduler": "karras"}),
     ("DPM++ SDE Karras", "sample_dpmpp_sde", {"scheduler": "karras"}),
 ]
 
+# samplers_diffusers = [
+#     ("DDIMScheduler", "diffusers.schedulers.DDIMScheduler", {})
+#     ("DDPMScheduler", "diffusers.schedulers.DDPMScheduler", {})
+#     ("DEISMultistepScheduler", "diffusers.schedulers.DEISMultistepScheduler", {})
+# ]
+
 start_time = time.time()
 
 scheduler = DDIMScheduler.from_pretrained(
@@ -62,22 +66,22 @@ scheduler = DDIMScheduler.from_pretrained(
 )
 vae = AutoencoderKL.from_pretrained(
     "stabilityai/sd-vae-ft-ema", 
-    torch_dtype=torch.float32
+    torch_dtype=torch.float16
 )
 text_encoder = CLIPTextModel.from_pretrained(
     base_model,
     subfolder="text_encoder",
-    torch_dtype=torch.float32,
+    torch_dtype=torch.float16,
 )
 tokenizer = CLIPTokenizer.from_pretrained(
     base_model,
     subfolder="tokenizer",
-    torch_dtype=torch.float32,
+    torch_dtype=torch.float16,
 )
 unet = UNet2DConditionModel.from_pretrained(
     base_model,
     subfolder="unet",
-    torch_dtype=torch.float32,
+    torch_dtype=torch.float16,
 )
 pipe = StableDiffusionPipeline(
     text_encoder=text_encoder,
@@ -88,15 +92,21 @@ pipe = StableDiffusionPipeline(
 )
 
 unet.set_attn_processor(CrossAttnProcessor)
+pipe.set_clip_skip(clip_skip)
 if torch.cuda.is_available():
     pipe = pipe.to("cuda")
-    
-def get_model_list():
-    return models
-
 
-unet_cache = dict()
 
+def get_model_list():
+    model_available = []
+    for model in models:
+        if Path(model[1]).is_dir():
+            model_available.append(model)
+    return model_available
+
+unet_cache = {
+    base_name: unet
+}
 
 def get_model(name):
     keys = [k[0] for k in models]
@@ -107,15 +117,14 @@ def get_model(name):
             unet = UNet2DConditionModel.from_pretrained(
                 models[keys.index(name)][1],
                 subfolder="unet",
-                torch_dtype=torch.float32,
+                torch_dtype=torch.float16,
             )
             unet_cache[name] = unet
-    
+
     g_unet = unet_cache[name]
     g_unet.set_attn_processor(None)
     return g_unet
 
-
 def error_str(error, title="Error"):
     return (
         f"""#### {title}
@@ -132,7 +141,7 @@ def restore_all():
     global te_base_weight, tokenizer
     text_encoder.get_input_embeddings().weight.data = te_base_weight
     tokenizer = CLIPTokenizer.from_pretrained(
-        "/root/workspace/storage/models/orangemix",
+        base_model,
         subfolder="tokenizer",
         torch_dtype=torch.float16,
     )
@@ -163,11 +172,8 @@ def inference(
     global pipe, unet, tokenizer, text_encoder
     if seed is None or seed == 0:
         seed = random.randint(0, 2147483647)
-    if torch.cuda.is_available():
-        generator = torch.Generator("cuda").manual_seed(int(seed))
-    else:
-        generator = torch.Generator().manual_seed(int(seed))
-    
+    generator = torch.Generator("cuda").manual_seed(int(seed))
+
     local_unet = get_model(model)
     if lora_state is not None and lora_state != "":
         load_lora_attn_procs(lora_state, local_unet, lora_scale)
@@ -189,15 +195,16 @@ def inference(
                 loaded_learned_embeds = load_file(file, device="cpu")
             loaded_learned_embeds = loaded_learned_embeds["string_to_param"]["*"]
             added_length = tokenizer.add_tokens(name)
-            
+
             assert added_length == loaded_learned_embeds.shape[0]
             delta_weight.append(loaded_learned_embeds)
 
         delta_weight = torch.cat(delta_weight, dim=0)
         text_encoder.resize_token_embeddings(len(tokenizer))
-        text_encoder.get_input_embeddings().weight.data[-delta_weight.shape[0]:] = delta_weight
+        text_encoder.get_input_embeddings().weight.data[
+            -delta_weight.shape[0] :
+        ] = delta_weight
 
-    
     config = {
         "negative_prompt": neg_prompt,
         "num_inference_steps": int(steps),
@@ -275,6 +282,9 @@ def apply_new_res(w, h, state):
 
 
 def detect_text(text, state, width, height):
+    
+    if text is None or text == "":
+        return None, None, None, None
 
     t = text.split(",")
     new_state = {}
@@ -287,11 +297,13 @@ def detect_text(text, state, width, height):
             new_state[item] = {
                 "map": state[item]["map"],
                 "weight": state[item]["weight"],
+                "mask_outsides": state[item]["weight"],
             }
         else:
             new_state[item] = {
                 "map": None,
                 "weight": 0.5,
+                "mask_outsides": False
             }
     update = gr.Radio.update(choices=[key for key in new_state.keys()], value=None)
     update_img = gr.update(value=create_mixed_img("", new_state, width, height))
@@ -314,28 +326,43 @@ def resize(img, w, h):
 def switch_canvas(entry, state, width, height):
     if entry == None:
         return None, 0.5, create_mixed_img("", state, width, height)
+
     return (
         gr.update(value=None, interactive=True),
-        gr.update(value=state[entry]["weight"]),
+        gr.update(value=state[entry]["weight"] if entry in state else 0.5),
+        gr.update(value=state[entry]["mask_outsides"] if entry in state else False),
         create_mixed_img(entry, state, width, height),
     )
 
 
 def apply_canvas(selected, draw, state, w, h):
-    w, h = int(w), int(h)
-    state[selected]["map"] = resize(draw, w, h)
+    if selected in state:
+        w, h = int(w), int(h)
+        state[selected]["map"] = resize(draw, w, h)
     return state, gr.Image.update(value=create_mixed_img(selected, state, w, h))
 
 
 def apply_weight(selected, weight, state):
-    state[selected]["weight"] = weight
+    if selected in state:
+        state[selected]["weight"] = weight
+    return state
+
+
+def apply_option(selected, mask, state):
+    if selected in state:
+        state[selected]["mask_outsides"] = mask
     return state
 
 
 # sp2, radio, width, height, global_stats
-def apply_image(image, selected, w, h, strgength, state):
-    if selected is not None:
-        state[selected] = {"map": resize(image, w, h), "weight": strgength}
+def apply_image(image, selected, w, h, strgength, mask, state):
+    if selected in state:
+        state[selected] = {
+            "map": resize(image, w, h), 
+            "weight": strgength, 
+            "mask_outsides": mask
+        }
+        
     return state, gr.Image.update(value=create_mixed_img(selected, state, w, h))
 
 
@@ -356,11 +383,24 @@ def add_net(files, ti_state, lora_state):
         else:
             ti_state[stripedname] = file.name
 
-    return ti_state, lora_state, gr.Text.update(f"{[key for key in ti_state.keys()]}"), gr.Text.update(f"{lora_state}"), gr.Files.update(value=None)
+    return (
+        ti_state,
+        lora_state,
+        gr.Text.update(f"{[key for key in ti_state.keys()]}"),
+        gr.Text.update(f"{lora_state}"),
+        gr.Files.update(value=None),
+    )
+
 
 # [ti_state, lora_state, ti_vals, lora_vals, uploads]
 def clean_states(ti_state, lora_state):
-    return dict(), None, gr.Text.update(f""), gr.Text.update(f""), gr.File.update(value=None)
+    return (
+        dict(),
+        None,
+        gr.Text.update(f""),
+        gr.Text.update(f""),
+        gr.File.update(value=None),
+    )
 
 
 latent_upscale_modes = {
@@ -564,15 +604,15 @@ with gr.Blocks(css=css) as demo:
                     with gr.Row():
                         with gr.Column(scale=90):
                             ti_vals = gr.Text(label="Loaded embeddings")
-                        
+
                     with gr.Row():
                         with gr.Column(scale=90):
                             lora_vals = gr.Text(label="Loaded loras")
 
                 with gr.Row():
-                        
+
                     uploads = gr.Files(label="Upload new embeddings/lora")
-                    
+
                     with gr.Column():
                         lora_scale = gr.Slider(
                             label="Lora scale",
@@ -583,12 +623,16 @@ with gr.Blocks(css=css) as demo:
                         )
                         btn = gr.Button(value="Upload")
                         btn_del = gr.Button(value="Reset")
-                        
+
                 btn.click(
-                    add_net, inputs=[uploads, ti_state, lora_state], outputs=[ti_state, lora_state, ti_vals, lora_vals, uploads]
+                    add_net,
+                    inputs=[uploads, ti_state, lora_state],
+                    outputs=[ti_state, lora_state, ti_vals, lora_vals, uploads],
                 )
                 btn_del.click(
-                    clean_states, inputs=[ti_state, lora_state], outputs=[ti_state, lora_state, ti_vals, lora_vals, uploads]
+                    clean_states,
+                    inputs=[ti_state, lora_state],
+                    outputs=[ti_state, lora_state, ti_vals, lora_vals, uploads],
                 )
 
         # error_output = gr.Markdown()
@@ -653,6 +697,11 @@ with gr.Blocks(css=css) as demo:
                     interactive=False,
                 )
 
+                mask_outsides = gr.Checkbox(
+                    label="Mask other areas", 
+                    value=False
+                )
+
                 strength = gr.Slider(
                     label="Token strength",
                     minimum=0,
@@ -660,6 +709,7 @@ with gr.Blocks(css=css) as demo:
                     step=0.01,
                     value=0.5,
                 )
+                
 
                 sk_update.click(
                     detect_text,
@@ -669,7 +719,7 @@ with gr.Blocks(css=css) as demo:
                 radio.change(
                     switch_canvas,
                     inputs=[radio, global_stats, width, height],
-                    outputs=[sp, strength, rendered],
+                    outputs=[sp, strength, mask_outsides, rendered],
                 )
                 sp.edit(
                     apply_canvas,
@@ -681,6 +731,11 @@ with gr.Blocks(css=css) as demo:
                     inputs=[radio, strength, global_stats],
                     outputs=[global_stats],
                 )
+                mask_outsides.change(
+                    apply_option,
+                    inputs=[radio, mask_outsides, global_stats],
+                    outputs=[global_stats],
+                )
 
             with gr.Tab("UploadFile"):
 
@@ -689,6 +744,11 @@ with gr.Blocks(css=css) as demo:
                     source="upload",
                     shape=(512, 512),
                 )
+                
+                mask_outsides2 = gr.Checkbox(
+                    label="Mask other areas", 
+                    value=False
+                )
 
                 strength2 = gr.Slider(
                     label="Token strength",
@@ -701,7 +761,7 @@ with gr.Blocks(css=css) as demo:
                 apply_style = gr.Button(value="Apply")
                 apply_style.click(
                     apply_image,
-                    inputs=[sp2, radio, width, height, strength2, global_stats],
+                    inputs=[sp2, radio, width, height, strength2, mask_outsides2, global_stats],
                     outputs=[global_stats, rendered],
                 )
 
@@ -740,7 +800,7 @@ with gr.Blocks(css=css) as demo:
         ti_state,
         model,
         lora_state,
-        lora_scale
+        lora_scale,
     ]
     outputs = [image_out]
     prompt.submit(inference, inputs=inputs, outputs=outputs)