Update app.py

app.py

@@ -51,6 +51,7 @@ pipe = pipeline("text-classification", model=model_ckpt)
 HF_TOKEN = os.environ.get("HUGGING_FACE_HUB_TOKEN")
 
 device="cuda" if torch.cuda.is_available() else "cpu"
+pipe_scribble, pipe_depth, pipe_img2img = None, None, None
 
 hidden_booster_text = "masterpiece++, best quality++, ultra-detailed+ +, unity 8k wallpaper+, illustration+, anime style+, intricate, fluid simulation, sharp edges. glossy++, Smooth++, detailed eyes++, best quality++,4k++,8k++,highres++,masterpiece++,ultra- detailed,realistic++,photorealistic++,photo-realistic++,depth of field, ultra-high definition, highly detailed, natural lighting, sharp focus, cinematic, hyperrealism,extremely detailed"
 hidden_negative = "bad anatomy, disfigured, poorly drawn,deformed, mutation, malformation, deformed, mutated, disfigured, deformed eyes+, bad face++, bad hands, poorly drawn hands, malformed hands, extra arms++, extra legs++, Fused body+, Fused hands+, Fused legs+, missing arms, missing limb, extra digit+, fewer digits, floating limbs, disconnected limbs, inaccurate limb, bad fingers, missing fingers, ugly face, long body++"
@@ -64,54 +65,81 @@ def translate(prompt):
     tgt_text = ''.join(tgt_text)[:-1]
     return tgt_text
 
+     
 
-hed = HEDdetector.from_pretrained('lllyasviel/ControlNet')
+def load_pipe_scribble():
+    if pipe_scribble is None:
+        hed = HEDdetector.from_pretrained('lllyasviel/ControlNet')
 
-controlnet_scribble = ControlNetModel.from_pretrained(
-    "lllyasviel/sd-controlnet-scribble", torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False, )
-
-pipe_scribble = StableDiffusionControlNetPipeline.from_single_file(
-    "https://huggingface.co/shellypeng/anime-god/blob/main/animeGod_v10.safetensors", controlnet=controlnet_scribble, safety_checker=None, requires_safety_checker=False,
-    torch_dtype=torch.float16, token=HF_TOKEN
-)
-
-pipe_scribble.load_lora_weights("shellypeng/lora2")
-pipe_scribble.fuse_lora(lora_scale=0.1)
-
-pipe_scribble.load_textual_inversion("shellypeng/textinv1")
-pipe_scribble.load_textual_inversion("shellypeng/textinv2")
-pipe_scribble.load_textual_inversion("shellypeng/textinv3")
-pipe_scribble.load_textual_inversion("shellypeng/textinv4")
-pipe_scribble.scheduler = DPMSolverMultistepScheduler.from_config(pipe_scribble.scheduler.config, use_karras_sigmas=True)
-pipe_scribble.safety_checker = None
-pipe_scribble.requires_safety_checker = False
-pipe_scribble.to(device)
-pipe_scribble.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
-
-
-depth_estimator = pipeline('depth-estimation')
-
-controlnet_depth = ControlNetModel.from_pretrained(
-    "lllyasviel/sd-controlnet-depth", torch_dtype=torch.float16
-)
+        controlnet_scribble = ControlNetModel.from_pretrained(
+            "lllyasviel/sd-controlnet-scribble", torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False, )
+        
+        pipe_scribble = StableDiffusionControlNetPipeline.from_single_file(
+            "https://huggingface.co/shellypeng/anime-god/blob/main/animeGod_v10.safetensors", controlnet=controlnet_scribble, safety_checker=None, requires_safety_checker=False,
+            torch_dtype=torch.float16, token=HF_TOKEN
+        )
+        
+        pipe_scribble.load_lora_weights("shellypeng/lora2")
+        pipe_scribble.fuse_lora(lora_scale=0.1)
+        
+        pipe_scribble.load_textual_inversion("shellypeng/textinv1")
+        pipe_scribble.load_textual_inversion("shellypeng/textinv2")
+        pipe_scribble.load_textual_inversion("shellypeng/textinv3")
+        pipe_scribble.load_textual_inversion("shellypeng/textinv4")
+        pipe_scribble.scheduler = DPMSolverMultistepScheduler.from_config(pipe_scribble.scheduler.config, use_karras_sigmas=True)
+        pipe_scribble.safety_checker = None
+        pipe_scribble.requires_safety_checker = False
+        pipe_scribble.to(device)
+        pipe_scribble.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
+
+
+def load_pipe_depth():
+    if pipe_depth is None:
+        depth_estimator = pipeline('depth-estimation')
+
+        controlnet_depth = ControlNetModel.from_pretrained(
+            "lllyasviel/sd-controlnet-depth", torch_dtype=torch.float16
+        )
+        
+        
+        pipe_depth = StableDiffusionControlNetPipeline.from_single_file(
+            "https://huggingface.co/shellypeng/anime-god/blob/main/animeGod_v10.safetensors", controlnet=controlnet_depth,
+            torch_dtype=torch.float16,
+        )
+        pipe_depth.load_lora_weights("shellypeng/lora1")
+        pipe_depth.fuse_lora(lora_scale=1.5)
+        
+        pipe_depth.load_textual_inversion("shellypeng/textinv1")
+        pipe_depth.load_textual_inversion("shellypeng/textinv2")
+        pipe_depth.load_textual_inversion("shellypeng/textinv3")
+        pipe_depth.load_textual_inversion("shellypeng/textinv4")
+        pipe_depth.scheduler = DPMSolverMultistepScheduler.from_config(pipe_depth.scheduler.config, use_karras_sigmas=True)
+        def dummy(images, **kwargs):
+            return images, False
+        pipe_depth.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
+        pipe_depth.to(device)
+        
+def load_pipe_img2img():
+    if pipe_img2img is None:
+        pipe_img2img = StableDiffusionImg2ImgPipeline.from_single_file("https://huggingface.co/shellypeng/anime-god/blob/main/animeGod_v10.safetensors",
+                                                               torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False, token=HF_TOKEN)
 
+        pipe_img2img.load_lora_weights("shellypeng/lora1")
+        pipe_img2img.fuse_lora(lora_scale=0.1)
+        pipe_img2img.load_lora_weights("shellypeng/lora2", token=HF_TOKEN)
+        pipe_img2img.fuse_lora(lora_scale=0.1)
+        
+        pipe_img2img.load_textual_inversion("shellypeng/textinv1")
+        pipe_img2img.load_textual_inversion("shellypeng/textinv2")
+        pipe_img2img.load_textual_inversion("shellypeng/textinv3")
+        pipe_img2img.load_textual_inversion("shellypeng/textinv4")
+        pipe_img2img.scheduler = DPMSolverMultistepScheduler.from_config(pipe_img2img.scheduler.config, use_karras_sigmas=True)
+        pipe_img2img.safety_checker = None
+        pipe_img2img.requires_safety_checker = False
+        pipe_img2img.to(device)
+        
+        pipe_img2img.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
 
-pipe_depth = StableDiffusionControlNetPipeline.from_single_file(
-    "https://huggingface.co/shellypeng/anime-god/blob/main/animeGod_v10.safetensors", controlnet=controlnet_depth,
-    torch_dtype=torch.float16,
-)
-pipe_depth.load_lora_weights("shellypeng/lora1")
-pipe_depth.fuse_lora(lora_scale=1.5)
-
-pipe_depth.load_textual_inversion("shellypeng/textinv1")
-pipe_depth.load_textual_inversion("shellypeng/textinv2")
-pipe_depth.load_textual_inversion("shellypeng/textinv3")
-pipe_depth.load_textual_inversion("shellypeng/textinv4")
-pipe_depth.scheduler = DPMSolverMultistepScheduler.from_config(pipe_depth.scheduler.config, use_karras_sigmas=True)
-def dummy(images, **kwargs):
-    return images, False
-pipe_depth.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
-pipe_depth.to(device)
 
 def real_to_anime(text, input_img):
     """
@@ -119,6 +147,7 @@ def real_to_anime(text, input_img):
     include Adetailer, detail tweaker lora, prompt backend include: beautiful eyes, beautiful face, beautiful hand, (maybe infer from user's prompt for gesture and facial
     expression to improve hand)
     """
+    load_pipe_depth()
     input_img = Image.fromarray(input_img)
     input_img = load_image(input_img)
     input_img = depth_estimator(input_img)['depth']
@@ -138,7 +167,7 @@ def scribble_to_image(text, neg_prompt_box, input_img):
     include Adetailer, detail tweaker lora, prompt backend include: beautiful eyes, beautiful face, beautiful hand, (maybe infer from user's prompt for gesture and facial
     expression to improve hand)
     """
-
+    load_pipe_scribble()
 
 
 # if auto detect detects chinese => auto turn on chinese prompting checkbox
@@ -170,6 +199,7 @@ def real_img2img_to_anime(text, neg_prompt_box, input_img):
     include Adetailer, detail tweaker lora, prompt backend include: beautiful eyes, beautiful face, beautiful hand, (maybe infer from user's prompt for gesture and facial
     expression to improve hand)
     """
+    load_pipe_img2img()
     input_img = Image.fromarray(input_img)
     input_img = load_image(input_img)
     lang_check_label = pipe(text, top_k=1, truncation=True)[0]['label']
@@ -202,24 +232,6 @@ theme = gr.themes.Soft(
 )
 
 
-pipe_img2img = StableDiffusionImg2ImgPipeline.from_single_file("https://huggingface.co/shellypeng/anime-god/blob/main/animeGod_v10.safetensors",
-                                                               torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False, token=HF_TOKEN)
-
-pipe_img2img.load_lora_weights("shellypeng/lora1")
-pipe_img2img.fuse_lora(lora_scale=0.1)
-pipe_img2img.load_lora_weights("shellypeng/lora2", token=HF_TOKEN)
-pipe_img2img.fuse_lora(lora_scale=0.1)
-
-pipe_img2img.load_textual_inversion("shellypeng/textinv1")
-pipe_img2img.load_textual_inversion("shellypeng/textinv2")
-pipe_img2img.load_textual_inversion("shellypeng/textinv3")
-pipe_img2img.load_textual_inversion("shellypeng/textinv4")
-pipe_img2img.scheduler = DPMSolverMultistepScheduler.from_config(pipe_img2img.scheduler.config, use_karras_sigmas=True)
-pipe_img2img.safety_checker = None
-pipe_img2img.requires_safety_checker = False
-pipe_img2img.to(device)
-
-pipe_img2img.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
 
 def zh_prompt_info(text, neg_text, chinese_check):
     can_raise_info = ""
@@ -302,7 +314,7 @@ with gr.Blocks(theme=theme, css="footer {visibility: hidden}", title="ShellAI Ap
             ["Beautiful girl, smiling, bun, bun hair, black hair, beautiful eyes, black dress, elegant, red carpet photo","ugly, bad quality", "emma.jpg"]
         ]
         
-        gr.Examples(examples=example_img2img, inputs=[prompt_box, neg_prompt_box, image_box], outputs=[image1, image2, image3, image4], fn=mult_thread_img2img, cache_examples=True)
+        # gr.Examples(examples=example_img2img, inputs=[prompt_box, neg_prompt_box, image_box], outputs=[image1, image2, image3, image4], fn=mult_thread_img2img, cache_examples=True)
 
         gr.on(triggers=[prompt_box.submit, gen_btn.click],fn=mult_thread_lang_class, inputs=[prompt_box, neg_prompt_box, chinese_check], outputs=[chinese_check], show_progress=False)
         gr.on(triggers=[prompt_box.submit, gen_btn.click],fn=real_to_anime, inputs=[prompt_box, image_box], outputs=[image1, image2, image3, image4])
@@ -337,7 +349,7 @@ with gr.Blocks(theme=theme, css="footer {visibility: hidden}", title="ShellAI Ap
             ["Beautiful girl, smiling, bun, bun hair, black hair, beautiful eyes, black dress, elegant, red carpet photo","ugly, bad quality", "emma.jpg"]
         ]
         
-        gr.Examples(examples=example_img2img, inputs=[prompt_box, neg_prompt_box, image_box], outputs=[image1, image2, image3, image4], fn=mult_thread_img2img, cache_examples=True)
+        # gr.Examples(examples=example_img2img, inputs=[prompt_box, neg_prompt_box, image_box], outputs=[image1, image2, image3, image4], fn=mult_thread_img2img, cache_examples=True)
 
         gr.on(triggers=[prompt_box.submit, gen_btn.click],fn=mult_thread_lang_class, inputs=[prompt_box, neg_prompt_box, chinese_check], outputs=[chinese_check], show_progress=False)
         gr.on(triggers=[prompt_box.submit, gen_btn.click],fn=mult_thread_img2img, inputs=[prompt_box, neg_prompt_box, image_box], outputs=[image1, image2, image3, image4])
@@ -371,7 +383,7 @@ with gr.Blocks(theme=theme, css="footer {visibility: hidden}", title="ShellAI Ap
             ["a beautiful girl spreading her arms, blue hair, long hair, hat with flowers on its edge, smiling++, dynamic, black dress, park background, birds, trees, flowers, grass","ugly, worst quality", "girl_spread.jpg"]
         ]
         
-        gr.Examples(examples=example_scribble_live2img, inputs=[prompt_box, neg_prompt_box, image_box], outputs=[image1, image2, image3, image4], fn=mult_thread_live_scribble, cache_examples=True)
+        # gr.Examples(examples=example_scribble_live2img, inputs=[prompt_box, neg_prompt_box, image_box], outputs=[image1, image2, image3, image4], fn=mult_thread_live_scribble, cache_examples=True)
 
         gr.on(triggers=[prompt_box.submit, gen_btn.click],fn=mult_thread_lang_class, inputs=[prompt_box, neg_prompt_box, chinese_check], outputs=[chinese_check], show_progress=False)
         gr.on(triggers=[prompt_box.submit, gen_btn.click],fn=mult_thread_live_scribble, inputs=[prompt_box, neg_prompt_box, image_box], outputs=[image1, image2, image3, image4])
@@ -405,7 +417,7 @@ with gr.Blocks(theme=theme, css="footer {visibility: hidden}", title="ShellAI Ap
             ["a man wearing a chinese clothes, closed eyes, handsome face, dragon on the clothes, expressionless face, indifferent, chinese building background","poor quality", "chinese_man.jpg"]
         ]
         
-        gr.Examples(examples=example_scribble2img, inputs=[prompt_box, neg_prompt_box, image_box], outputs=[image1, image2, image3, image4], fn=mult_thread_scribble, cache_examples=True)
+        # gr.Examples(examples=example_scribble2img, inputs=[prompt_box, neg_prompt_box, image_box], outputs=[image1, image2, image3, image4], fn=mult_thread_scribble, cache_examples=True)
 
         gr.on(triggers=[prompt_box.submit, gen_btn.click],fn=mult_thread_lang_class, inputs=[prompt_box, neg_prompt_box, chinese_check], outputs=[chinese_check], show_progress=False)
         gr.on(triggers=[prompt_box.submit, gen_btn.click],fn=mult_thread_scribble, inputs=[prompt_box, neg_prompt_box, image_box], outputs=[image1, image2, image3, image4])