Move 'reloading model weight' out of @spaces.GPU to avoid hanging

.gitignore

@@ -5,3 +5,4 @@ gradio_cached_examples/*
 gradio_cached_examples/
 samples/*
 samples/
+.gradio/certificate.pem

adaface/adaface_wrapper.py

@@ -399,8 +399,9 @@ class AdaFaceWrapper(nn.Module):
     # ref_img_strength is used only in the img2img pipeline.
     def forward(self, noise, prompt, negative_prompt=None, 
                 placeholder_tokens_pos='append',
+                guidance_scale=6.0, 
                 do_neg_id_prompt_weight=0,
-                guidance_scale=6.0, out_image_count=4, 
+                out_image_count=4, 
                 ref_img_strength=0.8, generator=None, verbose=False):
         noise = noise.to(device=self.device, dtype=torch.float16)
 

app.py

@@ -19,7 +19,7 @@ from insightface.utils import face_align
 from PIL import Image
 import torch
 import argparse
-# From command line read command adaface_ckpt_path
+
 parser = argparse.ArgumentParser()
 parser.add_argument("--adaface_encoder_types", type=str, nargs="+", default=["consistentID", "arc2face"],
                     choices=["arc2face", "consistentID"], help="Type(s) of the ID2Ada prompt encoders")
@@ -29,7 +29,7 @@ parser.add_argument('--model_style_type', type=str, default='realistic',
                     choices=["realistic", "anime", "photorealistic"], help="Type of the base model")
 parser.add_argument("--guidance_scale", type=float, default=8.0,
                     help="The guidance scale for the diffusion model. Default: 8.0")
-parser.add_argument("--do_neg_id_prompt_weight", type=float, default=0.2,
+parser.add_argument("--do_neg_id_prompt_weight", type=float, default=0,
                     help="The weight of added ID prompt embeddings into the negative prompt. Default: 0, disabled.")
 
 parser.add_argument('--gpu', type=int, default=None)
@@ -79,22 +79,12 @@ def get_clicked_image(data: gr.SelectData):
     return data.index
     
 @spaces.GPU
-def gen_init_images(uploaded_image_paths, model_style_type, prompt, out_image_count=3):
+def gen_init_images(uploaded_image_paths, prompt, 
+                    guidance_scale, do_neg_id_prompt_weight, out_image_count=3):
     global adaface, id_animator
     if uploaded_image_paths is None:
         print("No image uploaded")
         return None, None, None
-    
-    model_style_type = model_style_type.lower()
-    base_model_path = model_style_type2base_model_path[model_style_type]
-    # If the base model type is changed, reload the model.
-    if model_style_type != args.model_style_type:
-        id_animator = load_model(model_style_type=model_style_type, device=device)
-        adaface = AdaFaceWrapper(pipeline_name="text2img", base_model_path=base_model_path,
-                                 adaface_encoder_types=args.adaface_encoder_types,
-                                 adaface_ckpt_paths=[args.adaface_ckpt_path], device=device)
-        # Update base model type.
-        args.model_style_type = model_style_type
 
     # uploaded_image_paths is a list of tuples:
     # [('/tmp/gradio/249981e66a7c665aaaf1c7eaeb24949af4366c88/jensen huang.jpg', None)]
@@ -120,8 +110,9 @@ def gen_init_images(uploaded_image_paths, model_style_type, prompt, out_image_co
 
     # samples: A list of PIL Image instances.
     with torch.no_grad():
-        samples = adaface(noise, prompt, 
-                          placeholder_tokens_pos='append',
+        samples = adaface(noise, prompt, placeholder_tokens_pos='append',
+                          guidance_scale=guidance_scale, 
+                          do_neg_id_prompt_weight=do_neg_id_prompt_weight,
                           out_image_count=out_image_count, verbose=True)
 
     face_paths = []
@@ -137,23 +128,12 @@ def gen_init_images(uploaded_image_paths, model_style_type, prompt, out_image_co
 
 @spaces.GPU(duration=90)
 def generate_video(image_container, uploaded_image_paths, init_img_file_paths, init_img_selected_idx, 
-                   init_image_strength, init_image_final_weight, model_style_type,
+                   init_image_strength, init_image_final_weight,
                    prompt, negative_prompt, num_steps, video_length, guidance_scale, do_neg_id_prompt_weight, 
                    seed, attn_scale, image_embed_cfg_begin_scale, image_embed_cfg_end_scale,
                    is_adaface_enabled, adaface_ckpt_path, adaface_power_scale, 
                    id_animator_anneal_steps, progress=gr.Progress(track_tqdm=True)):
-
     global adaface, id_animator
-    model_style_type = model_style_type.lower()
-    base_model_path = model_style_type2base_model_path[model_style_type]
-    # If the base model type is changed, reload the model.
-    if model_style_type != args.model_style_type:
-        id_animator = load_model(model_style_type=model_style_type, device=device)
-        adaface = AdaFaceWrapper(pipeline_name="text2img", base_model_path=base_model_path,
-                                 adaface_encoder_types=args.adaface_encoder_types,
-                                 adaface_ckpt_paths=[args.adaface_ckpt_path], device=device)
-        # Update base model type.
-        args.model_style_type = model_style_type
 
     if prompt is None:
         prompt = ""
@@ -231,11 +211,24 @@ def generate_video(image_container, uploaded_image_paths, init_img_file_paths, i
     save_videos_grid(sample, save_sample_path)
     return save_sample_path
 
-def validate_prompt(prompt):
+def check_prompt_and_model_type(prompt, model_style_type):
+    global adaface, id_animator
+
+    model_style_type = model_style_type.lower()
+    base_model_path = model_style_type2base_model_path[model_style_type]
+    # If the base model type is changed, reload the model.
+    if model_style_type != args.model_style_type:
+        id_animator = load_model(model_style_type=model_style_type, device=device)
+        adaface = AdaFaceWrapper(pipeline_name="text2img", base_model_path=base_model_path,
+                                 adaface_encoder_types=args.adaface_encoder_types,
+                                 adaface_ckpt_paths=[args.adaface_ckpt_path], device=device)
+        # Update base model type.
+        args.model_style_type = model_style_type
+
     if not prompt:
         raise gr.Error("Prompt cannot be blank")
 
-with gr.Blocks(css=css) as demo:
+with gr.Blocks(css=css, theme=gr.themes.Origin()) as demo:
     gr.Markdown(
         """
         # AdaFace-Animate: Zero-Shot Subject-Driven Video Generation for Humans
@@ -269,7 +262,7 @@ with gr.Blocks(css=css) as demo:
                         file_count="multiple"
                     )
             image_container = gr.Image(label="image container", sources="upload", type="numpy", height=256, visible=False)
-            uploaded_files_gallery = gr.Gallery(label="Subject images", visible=False, columns=3, rows=1, height=300)
+            uploaded_files_gallery = gr.Gallery(label="Subject images", visible=False, columns=3, rows=2, height=300)
             with gr.Column(visible=False) as clear_button_column:
                 remove_and_reupload = gr.ClearButton(value="Remove and upload subject images", components=files, size="sm")
 
@@ -348,6 +341,7 @@ with gr.Blocks(css=css) as demo:
                 maximum=0.9,
                 step=0.1,
                 value=args.do_neg_id_prompt_weight,
+                visible=False
             )
 
             seed = gr.Slider(
@@ -447,15 +441,18 @@ with gr.Blocks(css=css) as demo:
         files.upload(fn=swap_to_gallery, inputs=files,     outputs=[uploaded_files_gallery, clear_button_column, files])
         remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files_gallery, clear_button_column, files, init_img_selected_idx])
 
-        init_img_files.upload(fn=swap_to_gallery, inputs=init_img_files, outputs=[uploaded_init_img_gallery, init_clear_button_column, init_img_files])
-        remove_init_and_reupload.click(fn=remove_back_to_files,        outputs=[uploaded_init_img_gallery, init_clear_button_column, 
-                                                                                init_img_files, init_img_selected_idx])
-        gen_init.click(fn=gen_init_images, inputs=[uploaded_files_gallery, model_style_type, prompt], 
+        init_img_files.upload(fn=swap_to_gallery, inputs=init_img_files, 
+                              outputs=[uploaded_init_img_gallery, init_clear_button_column, init_img_files])
+        remove_init_and_reupload.click(fn=remove_back_to_files, 
+                                       outputs=[uploaded_init_img_gallery, init_clear_button_column, 
+                                                init_img_files, init_img_selected_idx])
+        gen_init.click(fn=gen_init_images, inputs=[uploaded_files_gallery, prompt, 
+                                                   guidance_scale, do_neg_id_prompt_weight], 
                        outputs=[uploaded_init_img_gallery, init_img_files, init_clear_button_column])
         uploaded_init_img_gallery.select(fn=get_clicked_image, inputs=None, outputs=init_img_selected_idx)
 
-        submit.click(fn=validate_prompt,
-                     inputs=[prompt],outputs=None).success(
+        submit.click(fn=check_prompt_and_model_type,
+                     inputs=[prompt, model_style_type],outputs=None).success(
             fn=randomize_seed_fn,
             inputs=[seed, randomize_seed],
             outputs=seed,
@@ -463,8 +460,8 @@ with gr.Blocks(css=css) as demo:
             api_name=False,
         ).then(
                  fn=generate_video,
-                 inputs=[image_container, files, init_img_files, init_img_selected_idx, init_image_strength, 
-                         init_image_final_weight, model_style_type,
+                 inputs=[image_container, files, 
+                         init_img_files, init_img_selected_idx, init_image_strength, init_image_final_weight,
                          prompt, negative_prompt, num_steps, video_length, guidance_scale, do_neg_id_prompt_weight,
                          seed, attn_scale, image_embed_cfg_begin_scale, image_embed_cfg_end_scale,
                          is_adaface_enabled, adaface_ckpt_path, adaface_power_scale, id_animator_anneal_steps],

models/adaface/{VGGface2_HQ_masks2024-10-08T14-42-05_zero3-ada-30000.pt → VGGface2_HQ_masks2024-10-13T11-21-07_zero3-ada-9000.pt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:34bbdaa97fb2da9e2aae4204bfd2f5c1565a84c664520a5f537129419ecb53fa
-size 1814921594
+oid sha256:c8d4dae6291aa783c03de6e54997a79b7d61f7bdc2457afb6303bcfc290d0956
+size 1814921242