0516_fix_errors

app.py

@@ -10,6 +10,7 @@ import gradio as gr
 import requests
 import time
 import random
+from style_template import styles
 import numpy as np
 import torch
 import os
@@ -26,6 +27,7 @@ from models.celeb_embeddings import embedding_forward
 import models.embedding_manager
 import importlib
 import time
+
 import os
 # os.environ['GRADIO_TEMP_DIR'] = 'qinghewang/tmp'
 
@@ -128,30 +130,13 @@ woman_Embedding_Manager = models.embedding_manager.EmbeddingManagerId_adain(
         loss_type = embedding_manager_config.model.personalization_config.params.loss_type,
         vit_out_dim = input_dim,
 )
+
 text_encoder.text_model.embeddings.forward = original_forward
 
+STYLE_NAMES = list(styles.keys())
 DEFAULT_STYLE_NAME = "Watercolor"
 MAX_SEED = np.iinfo(np.int32).max
 
-def remove_tips():
-    return gr.update(visible=False)
-
-def response(choice, gender_GAN):
-    c = ""
-    e = ""
-    if choice == "Create a new character":
-        c = "create"
-    elif choice == "Still use this character":
-        c = "continue"
-    
-    if gender_GAN == "Normal":
-        e = "normal_GAN"
-    elif gender_GAN == "Man":
-        e = "man_GAN"
-    elif gender_GAN == "Woman":
-        e = "woman_GAN"
-    
-    return c, e
 
 def replace_phrases(prompt):
     replacements = {
@@ -174,47 +159,42 @@ def handle_prompts(prompts_array):
     
 
 @spaces.GPU
-def generate_image(experiment_name, label, prompts_array, chose_emb):
+def generate_image(chose_emb, choice, gender_GAN, prompts_array):
     prompts = handle_prompts(prompts_array)
     
-    print("experiment_name:",experiment_name)
+    print("gender:",gender_GAN)
 
-    if experiment_name == "normal_GAN":
+    if gender_GAN == "Normal":
         steps = 10000
         Embedding_Manager = normal_Embedding_Manager
-    elif experiment_name == "man_GAN":
+    elif gender_GAN == "Man":
         steps = 7000
         Embedding_Manager = man_Embedding_Manager
-    elif experiment_name == "woman_GAN":
+    elif gender_GAN == "Woman":
         steps = 6000
         Embedding_Manager = woman_Embedding_Manager
     else:
         print("Hello, please notice this ^_^")
         assert 0
     
-    embedding_path = os.path.join("training_weight", experiment_name, "embeddings_manager-{}.pt".format(str(steps)))
+    embedding_path = os.path.join("training_weight", gender_GAN, "embeddings_manager-{}.pt".format(str(steps)))
     Embedding_Manager.load(embedding_path)
     print("embedding_path:",embedding_path)
-    print("label:",label)
+    print("choice:",choice)
     
-    index = "0"
-    save_dir = os.path.join("test_results/" + experiment_name, index)
+    # index = "0"
+    save_dir = os.path.join("test_results/" + gender_GAN)   # , index
     os.makedirs(save_dir, exist_ok=True)
-    ran_emb_path = os.path.join(save_dir, "ran_embeddings.pt")
-    test_emb_path = os.path.join(save_dir, "id_embeddings.pt")
     
     random_embedding = torch.randn(1, 1, input_dim).to(device)
-    if label == "create":
-        print("new")
-        torch.save(random_embedding, ran_emb_path)
+    if choice == "Create a new character":
         _, emb_dict = Embedding_Manager(tokenized_text=None, embedded_text=None, name_batch=None, random_embeddings = random_embedding, timesteps = None,)
-        # text_encoder.text_model.embeddings.forward = original_forward
         test_emb = emb_dict["adained_total_embedding"].to(device)
-        torch.save(test_emb, test_emb_path)
-    elif label == "continue":
-        print("old")
+    elif choice == "Still use this character":
         test_emb = torch.load(chose_emb).cuda()
-        # text_encoder.text_model.embeddings.forward = original_forward
+    
+    test_emb_path = os.path.join(save_dir, "id_embeddings.pt")
+    torch.save(test_emb, test_emb_path)
     
     v1_emb = test_emb[:, 0]
     v2_emb = test_emb[:, 1]
@@ -229,107 +209,57 @@ def generate_image(experiment_name, label, prompts_array, chose_emb):
         text_encoder.get_input_embeddings().weight.data[token_id] = embedding
     
     total_results = []
+    i = 0
     for prompt in prompts:
         image = pipe(prompt, guidance_scale = 8.5).images
         total_results = image + total_results
-        yield total_results, test_emb_path
+        i+=1
+        if i < len(prompts):
+            yield total_results, gr.update(visible=True, value="<h3>(Not Finished) Generating ···</h3>"), test_emb_path
+        else:
+            yield total_results, gr.update(visible=True, value="<h3>Generation Finished</h3>"), test_emb_path
 
 def get_example():
     case = [
         [
             'demo_embeddings/example_1.pt',
+            'Still use this character',
             "Normal",
-            "Still use this character",
             "a photo of a person\na person as a small child\na person as a 20 years old person\na person as a 80 years old person\na person reading a book\na person in the sunset\n",
         ],
         [
             'demo_embeddings/example_2.pt',
+            'Still use this character',
             "Man",
-            "Still use this character",
             "a photo of a person\na person with a mustache and a hat\na person wearing headphoneswith red hair\na person with his dog\n",
         ],
         [
             'demo_embeddings/example_3.pt',
+            'Still use this character',
             "Woman",
-            "Still use this character",
             "a photo of a person\na person at a beach\na person as a police officer\na person wearing a birthday hat\n",
         ],
         [
             'demo_embeddings/example_4.pt',
+            'Still use this character',
             "Man",
-            "Still use this character",
             "a photo of a person\na person holding a bunch of flowers\na person in a lab coat\na person speaking at a podium\n",
         ],
         [
             'demo_embeddings/example_5.pt',
+            'Still use this character',
             "Woman",
-            "Still use this character",
             "a photo of a person\na person wearing a kimono\na person in Van Gogh style\nEthereal fantasy concept art of a person\n",
         ],
         [
             'demo_embeddings/example_6.pt',
+            'Still use this character',
             "Man",
-            "Still use this character",
             "a photo of a person\na person in the rain\na person meditating\na pencil sketch of a person\n",
         ],
     ]
     return case
 
-@spaces.GPU
-def run_for_examples(example_emb, gender_GAN, choice, prompts_array):
-    prompts = handle_prompts(prompts_array)
-    label, experiment_name = response(choice, gender_GAN)
-    if experiment_name == "normal_GAN":
-        steps = 10000
-        Embedding_Manager = normal_Embedding_Manager
-    elif experiment_name == "man_GAN":
-        steps = 7000
-        Embedding_Manager = man_Embedding_Manager
-    elif experiment_name == "woman_GAN":
-        steps = 6000
-        Embedding_Manager = woman_Embedding_Manager
-    else:
-        print("Hello, please notice this ^_^")
-        assert 0
-        
-    embedding_path = os.path.join("training_weight", experiment_name, "embeddings_manager-{}.pt".format(str(steps)))
-    Embedding_Manager.load(embedding_path)
-    print("embedding_path:",embedding_path)
-    print("label:",label)
-    
-    test_emb = torch.load(example_emb).cuda()
-    v1_emb = test_emb[:, 0]
-    v2_emb = test_emb[:, 1]
-    embeddings = [v1_emb, v2_emb]
-    
-    tokens = ["v1*", "v2*"]
-    tokenizer.add_tokens(tokens)
-    token_ids = tokenizer.convert_tokens_to_ids(tokens)
-    
-    text_encoder.resize_token_embeddings(len(tokenizer), pad_to_multiple_of = 8)
-    for token_id, embedding in zip(token_ids, embeddings):
-        text_encoder.get_input_embeddings().weight.data[token_id] = embedding
-    
-    total_results = []
-    i = 0
-    for prompt in prompts:
-        image = pipe(prompt, guidance_scale = 8.5).images
-        total_results = image + total_results
-        i+=1
-        if i < len(prompts):
-            yield total_results, gr.update(visible=True, value="<h3>(Not Finished) Generating ···</h3>")
-        else:
-            yield total_results, gr.update(visible=True, value="<h3>Generation Finished</h3>")
-
-
-def set_text_unfinished():
-    return gr.update(visible=True, value="<h3>(Not Finished) Generating ···</h3>")
- 
-def set_text_finished():
-    return gr.update(visible=True, value="<h3>Generation Finished</h3>")
-
-
-
 
 with gr.Blocks(css=css) as demo:    # css=css
     # binary_matrixes = gr.State([])
@@ -344,17 +274,12 @@ with gr.Blocks(css=css) as demo:    # css=css
             prompts_array = gr.Textbox(lines = 3, 
                     label="Prompts (each line corresponds to a frame).", 
                     info="Give simple prompt is enough to achieve good face fidelity",
-                    # placeholder="A photo of a person",
                     value="a photo of a person\na person reading a book\na person wearing a Christmas hat\na Fauvism painting of a person\n",
                     interactive=True)
             choice = gr.Radio(choices=["Create a new character", "Still use this character"], label="Choose your action")
 
-            gender_GAN = gr.Radio(choices=["Normal", "Man", "Woman"], label="Choose your model version")
-            
-            label = gr.Text(label="Select the action you want to take", visible=False)
-            experiment_name = gr.Text(label="Select the GAN you want to take", visible=False)
+            gender_GAN = gr.Radio(choices=["Normal", "Man", "Woman"], label="Choose your model version (Only work for 'Create a new character')")   # , disabled=False
             chose_emb = gr.File(label="Uploaded files", type="filepath", visible=False)
-            example_emb = gr.File(label="Uploaded files", type="filepath", visible=False)
             
             generate = gr.Button("Generate!😊", variant="primary")
 
@@ -363,33 +288,21 @@ with gr.Blocks(css=css) as demo:    # css=css
             generated_information = gr.Markdown(label="Generation Details", value="",visible=False)
             
         generate.click(
-            fn=set_text_unfinished, 
-            outputs=generated_information
-        ).then(
-            fn=response,
-            inputs=[choice, gender_GAN],
-            outputs=[label, experiment_name],    
-        ).then(
             fn=generate_image,
-            inputs=[experiment_name, label, prompts_array, chose_emb],
-            outputs=[gallery, chose_emb] 
-        ).then(
-            fn=set_text_finished,
-            outputs=generated_information
+            inputs=[chose_emb, choice, gender_GAN, prompts_array],
+            outputs=[gallery, generated_information, chose_emb] 
         )
         
+        
 
     gr.Examples(
         examples=get_example(),
-        inputs=[example_emb, gender_GAN, choice, prompts_array],
-        run_on_click=True,
-        fn=run_for_examples,
-        outputs=[gallery, generated_information],
+        inputs=[chose_emb, choice, gender_GAN, prompts_array],
+        run_on_click=False,
+        fn=generate_image,
+        outputs=[gallery, generated_information, chose_emb],
     )
     
     gr.Markdown(article)
-# demo.launch(server_name="0.0.0.0", share = False)
-# share_link = demo.launch(share=True)
-# print("Share this link: ", share_link)
 
 demo.launch()   # share=True