Update app.py

app.py

@@ -1,11 +1,81 @@
-from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, DPMSolverMultistepScheduler
+from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, DPMSolverMultistepScheduler, AutoencoderKL
 import gradio as gr
 import torch
 from PIL import Image
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+
+def convert_safetensors_to_bin(pipeline, state_dict, alpha = 0.4):
+    LORA_PREFIX_UNET = 'lora_unet'
+    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
+
+    visited = []
+
+    # directly update weight in diffusers model
+    for key in state_dict:
+        # it is suggested to print out the key, it usually will be something like below
+        # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight"
+
+        # as we have set the alpha beforehand, so just skip
+        if '.alpha' in key or key in visited:
+            continue
+
+        if 'text' in key:
+            layer_infos = key.split('.')[0].split(LORA_PREFIX_TEXT_ENCODER + '_')[-1].split('_')
+            curr_layer = pipeline.text_encoder
+        else:
+            layer_infos = key.split('.')[0].split(LORA_PREFIX_UNET + '_')[-1].split('_')
+            curr_layer = pipeline.unet
+
+        # find the target layer
+        temp_name = layer_infos.pop(0)
+        while len(layer_infos) > -1:
+            try:
+                curr_layer = curr_layer.__getattr__(temp_name)
+                if len(layer_infos) > 0:
+                    temp_name = layer_infos.pop(0)
+                elif len(layer_infos) == 0:
+                    break
+            except Exception:
+                if len(temp_name) > 0:
+                    temp_name += '_' + layer_infos.pop(0)
+                else:
+                    temp_name = layer_infos.pop(0)
+
+        # org_forward(x) + lora_up(lora_down(x)) * multiplier
+        pair_keys = []
+        if 'lora_down' in key:
+            pair_keys.append(key.replace('lora_down', 'lora_up'))
+            pair_keys.append(key)
+        else:
+            pair_keys.append(key)
+            pair_keys.append(key.replace('lora_up', 'lora_down'))
+
+        # update weight
+        if len(state_dict[pair_keys[0]].shape) == 4:
+            weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32)
+            weight_down = state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32)
+            curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3)
+        else:
+            weight_up = state_dict[pair_keys[0]].to(torch.float32)
+            weight_down = state_dict[pair_keys[1]].to(torch.float32)
+            curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down)
+
+        # update visited list
+        for item in pair_keys:
+            visited.append(item)
+
+    return pipeline
+
 
 model_id = 'andite/anything-v4.0'
 prefix = ''
-     
+lora_path = hf_hub_download(
+    "showee/showee-lora-v1.0", "showee-any4.0.safetensors"
+)
+vae_path = "./anything-v4.0-vae/diffusion_pytorch_model.bin"
+
 scheduler = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler")
 
 pipe = StableDiffusionPipeline.from_pretrained(
@@ -13,11 +83,23 @@ pipe = StableDiffusionPipeline.from_pretrained(
   torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
   scheduler=scheduler)
 
+pipe.vae.load_state_dict(torch.load(vae_path))
+
+state_dict = load_file(lora_path)
+pipe = convert_safetensors_to_bin(pipe, state_dict, 0.3)
+
+
 pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(
   model_id,
   torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
   scheduler=scheduler)
 
+pipe_i2i.vae.load_state_dict(torch.load(vae_path))
+
+state_dict_i2i = load_file(lora_path)
+pipe_i2i = convert_safetensors_to_bin(pipe, state_dict_i2i, 0.3)
+
+
 if torch.cuda.is_available():
   pipe = pipe.to("cuda")
   pipe_i2i = pipe_i2i.to("cuda")
@@ -76,14 +158,14 @@ with gr.Blocks(css=css) as demo:
         f"""
             <div class="main-div">
               <div>
-                <h1>Anything V4.0</h1>
+                <h1>Showee V1.0</h1>
               </div>
               <p>
-               Demo for <a href="https://huggingface.co/andite/anything-v4.0">Anything V4.0</a> Stable Diffusion model.<br>
+               Demo for <a href="https://huggingface.co/showee/showee-lora-v1.0">Showee V1.0</a> LoRA adaption weights fine-tuned from <a href="https://huggingface.co/andite/anything-v4.0">Anything V4.0</a> Stable Diffusion model.<br>
                {"Add the following tokens to your prompts for the model to work properly: <b>prefix</b>" if prefix else ""}
               </p>
-              Running on {"<b>GPU 🔥</b>" if torch.cuda.is_available() else f"<b>CPU 🥶</b>. For faster inference it is recommended to <b>upgrade to GPU in <a href='https://huggingface.co/spaces/akhaliq/anything-v4.0/settings'>Settings</a></b>"} after duplicating the space<br><br>
-              <a style="display:inline-block" href="https://huggingface.co/spaces/akhaliq/anything-v4.0?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
+              Running on {"<b>GPU 🔥</b>" if torch.cuda.is_available() else f"<b>CPU 🥶</b>. For faster inference it is recommended to <b>upgrade to GPU in <a href='https://huggingface.co/spaces/showee/showee-v1.0/settings'>Settings</a></b>"} after duplicating the space<br><br>
+              <a style="display:inline-block" href="https://huggingface.co/spaces/showee/showee-v1.0?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
             </div>
         """
     )
@@ -101,7 +183,15 @@ with gr.Blocks(css=css) as demo:
         with gr.Column(scale=45):
           with gr.Tab("Options"):
             with gr.Group():
-              neg_prompt = gr.Textbox(label="Negative prompt", placeholder="What to exclude from the image")
+              neg_prompt = gr.Textbox(label="Negative prompt",
+                                      placeholder="What to exclude from the image",
+                                      value="NSFW, lowres, ((bad anatomy)), ((bad hands)), text, missing finger, "
+                                            "extra digits, fewer digits, blurry, ((mutated hands and fingers)), "
+                                            "(poorly drawn face), ((mutation)), ((deformed face)), (ugly), "
+                                            "((bad proportions)), ((extra limbs)), extra face, (double head), "
+                                            "(extra head), ((extra feet)), monster, logo, cropped, worst quality, "
+                                            "low quality, normal quality, jpeg, humpbacked, long body, long neck, "
+                                            "((jpeg artifacts))")
               auto_prefix = gr.Checkbox(label="Prefix styling tokens automatically ()", value=prefix, visible=prefix)
 
               with gr.Row():
@@ -119,6 +209,25 @@ with gr.Blocks(css=css) as demo:
                 image = gr.Image(label="Image", height=256, tool="editor", type="pil")
                 strength = gr.Slider(label="Transformation strength", minimum=0, maximum=1, step=0.01, value=0.5)
 
+    gr.Examples(
+        [[
+            "masterpiece, best quality, ultra-detailed, illustration, portrait, 1girl, solo, white hair, green eyes, "
+            "aqua_eyes, cat_ears, :3, ahoge, dress, red_jacket, long_sleeves, bangs, black_legwear, hair_ornament, "
+            "hairclip",  8, 25, 768, 1024, 909198616
+        ],
+            [
+                "masterpiece, best quality, ultra-detailed, illustration, portrait, 1girl, :3, animal_ears, aqua_eyes, ahoge, "
+                "asymmetrical_legwear, bangs, black_footwear, black_skirt, breasts, cleavage, hair_ornament, hairclip, "
+                "long_hair, navel, thighhighs, smile", 7.5, 25, 512, 768, 9
+            ],
+            [
+                "masterpiece, best quality, ultra-detailed, illustration, portrait, 1girl, :3, animal_ears, aqua_eyes,  ahoge, seaside,"
+                "asymmetrical_legwear, bangs, black_footwear, black_skirt, breasts, cleavage, hair_ornament, hairclip, "
+                "long_hair, navel, thighhighs", 7.5, 25, 512, 512, 353573117
+            ]],
+        [prompt, guidance, steps, width, height, seed],
+    )
+
     auto_prefix.change(lambda x: gr.update(placeholder=f"{prefix} [your prompt]" if x else "[Your prompt]"), inputs=auto_prefix, outputs=prompt, queue=False)
 
     inputs = [prompt, guidance, steps, width, height, seed, image, strength, neg_prompt, auto_prefix]