Update app.py

app.py

@@ -16,9 +16,9 @@ from diffusers import StableDiffusionXLPipeline, DiffusionPipeline
 import anthropic
 import base64
 from datasets import load_dataset
+from PIL import Image 
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(device)
 
 word_list_dataset = load_dataset("EPFL-VILAB/4m-wordlist", data_files="list.txt", use_auth_token=True)
 word_list = word_list_dataset["train"]['text']
@@ -50,7 +50,6 @@ else:
     pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", use_safetensors=True)
     pipe = pipe.to(device)
 
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
@@ -75,13 +74,10 @@ css="""
 }
 """
 
-
 if torch.cuda.is_available():
     power_device = "GPU"
 else:
     power_device = "CPU"
-
-from PIL import Image 
   
 comment_images = [
     "comment_images/15.png",
@@ -158,7 +154,6 @@ examples = {
         "examples/A person reaching fo_0.png",
         "examples/A person reaching fo_1.png",
         "examples/A person reaching fo_2.png",
-        #"examples/A person reaching fo_3.png",
         "examples/A person reaching fo_4.png",
         "examples/A person reaching fo_5.png",
         "examples/A person reaching fo_6.png",
@@ -170,7 +165,6 @@ examples = {
         "examples/Abandoned robot at t_0.png",
         "examples/Abandoned robot at t_1.png",
         "examples/Abandoned robot at t_2.png",
-        #"examples/Abandoned robot at t_3.png",
         "examples/Abandoned robot at t_4.png",
         "examples/Abandoned robot at t_5.png",
         "examples/Abandoned robot at t_6.png",
@@ -182,7 +176,6 @@ examples = {
         "examples/Cityscape during a t_0.png",
         "examples/Cityscape during a t_1.png",
         "examples/Cityscape during a t_2.png",
-        #"examples/Cityscape during a t_3.png",
         "examples/Cityscape during a t_4.png",
         "examples/Cityscape during a t_5.png",
         "examples/Cityscape during a t_6.png",
@@ -194,7 +187,6 @@ examples = {
         "examples/Human in a frame_0.png",
         "examples/Human in a frame_1.png",
         "examples/Human in a frame_2.png",
-        #"examples/Human in a frame_3.png",
         "examples/Human in a frame_4.png",
         "examples/Human in a frame_5.png",
         "examples/Human in a frame_6.png",
@@ -206,7 +198,6 @@ examples = {
         "examples/Inside an abondoned _0.png",
         "examples/Inside an abondoned _1.png",
         "examples/Inside an abondoned _2.png",
-        #"examples/Inside an abondoned _3.png",
         "examples/Inside an abondoned _4.png",
         "examples/Inside an abondoned _5.png",
         "examples/Inside an abondoned _6.png",
@@ -218,7 +209,6 @@ examples = {
         "examples/Lonely astronaut in _0.png",
         "examples/Lonely astronaut in _1.png",
         "examples/Lonely astronaut in _2.png",
-        #"examples/Lonely astronaut in _3.png",
         "examples/Lonely astronaut in _4.png",
         "examples/Lonely astronaut in _5.png",
         "examples/Lonely astronaut in _6.png",
@@ -230,7 +220,6 @@ examples = {
         "examples/Painting of a lady_0.png",
         "examples/Painting of a lady_1.png",
         "examples/Painting of a lady_2.png",
-        #"examples/Painting of a lady_3.png",
         "examples/Painting of a lady_4.png",
         "examples/Painting of a lady_5.png",
         "examples/Painting of a lady_6.png",
@@ -251,7 +240,6 @@ def submit_comment(comment):
         comments[comment_images[0]] = comment
         comment_images.append(comment_images[0])
         comment_images = comment_images[1:]
-
         image_index = (image_index + 1) % len(comment_images)
 
     elif comment_images[0] in comments:
@@ -310,7 +298,7 @@ def clear_comments():
 def extract_vp_from_vpe():
     global comments
     if len(comments) < 8:
-        gr.Warning("Number of comments less than 8 could lead to errors.")
+        gr.Warning("Fewer than 8 comments may lead to errors.")
         
     prompt = """I will provide a set of artworks along with accompanying comments from a person. Analyze these artworks and the comments on them and identify artistic features such as present or mentioned colors, style, composition, mood, medium, texture, brushwork, lighting, shadow effects, perspective, and other noteworthy elements.
 Your task is to extract the artistic features the person likes and dislikes based on both the artworks' features and the person's comments. Focus solely on artistic aspects and refrain from considering subject matter.
@@ -351,8 +339,13 @@ Here are the images and their corresponding comments:
     generated_ids = vpe_model.generate(**inputs, max_new_tokens=2000, repetition_penalty=0.99, do_sample=False)
     generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
     
-    gr.Info("Visual preference successfully extracted.")
-    positive_vp, negative_vp = re.search('.* \nAssistant: Liked Art Features: (.*)\nDisliked Art Features: (.*)', generated_texts).groups()
+    if re.match(r"(.|\n)*Assistant: Liked Art Features: (.|\n)*Disliked Art Features: (.|\n)*", generated_texts):
+        positive_vp, negative_vp = re.search('.* \nAssistant: Liked Art Features: (.*)\nDisliked Art Features: (.*)', generated_texts).groups()
+        gr.Info("Visual preference successfully extracted.")
+    else:
+        positive_vp = ""
+        negative_vp = ""
+        gr.Warning("VP extraction failed: Please comment on more images.")
         
     return positive_vp, negative_vp
 
@@ -361,7 +354,6 @@ def extract_vp():
     if valid_api == "":
         positive_vp, negative_vp = extract_vp_from_vpe()
         
-
     else:
         client = anthropic.Anthropic(
             api_key=valid_api,
@@ -369,31 +361,23 @@ def extract_vp():
 
         prompt = """**Objective:**
 Analyze a set of artworks and accompanying comments from a person to identify artistic features they like and dislike.
-
 **Steps:**
-
 1. **Analyze Artworks and Comments:**
    - Examine each artwork for artistic features such as colors, style, composition, mood, medium, texture, brushwork, lighting, shadow effects, perspective, and other noteworthy elements.
    - Review the accompanying comments to understand the person's preferences and opinions on these features.
-
 2. **Identify Preferences:**
    - Extract artistic features that the person likes and dislikes based on the artworks' features and the comments.
    - Focus solely on artistic aspects and ignore the subject matter.
    - Convert the art features mentioned in the comments to well-known synonyms if needed.
-
 3. **Resolve Ambiguous Preferences:**
    - If the person expresses a preference without clearly stating its category (e.g., "I like the style" without specifying which style), identify these specific features from the images directly.
    - Make the person's preference understandable and independednt of the artworks.
-
 4. **Output Format:**
    - Create two concise lists of keywords: one for features the person likes and another for features they dislike.
    - Ensure the lists are in keyword format, divided by commas, without using sentences.
    - Maintain detail and accuracy for all comments and images.
-
 **Your Task:**
-
 Follow the example format and ensure that your output consists of two lists of keywords summarizing the person's preferences based on the artworks and comments provided. Consider all comments and images comprehensively.
-
 **Example**: example START:
 """
         messages = [
@@ -448,8 +432,13 @@ Follow the example format and ensure that your output consists of two lists of k
         )
 
         generated_text = message.content[0].text
-        positive_vp, negative_vp = re.search('.*Like.*:\n(.*)\n*Dislike.*:\n(.*)', generated_text).groups()
-    
+        if re.match(r"(.|\n)*Like.*:(.|\n)*Dislike.(.|\n)*", generated_texts):
+            positive_vp, negative_vp = re.search('.*Like.*:\n(.*)\n*Dislike.*:\n(.*)', generated_text).groups()
+            gr.Info("Visual preference successfully extracted.")
+        else:
+            positive_vp = ""
+            negative_vp = ""
+            gr.Warning("VP extraction failed: Please comment on more images.")
     
     return gr.Textbox(label="Liked visual attributes", lines=3, value=positive_vp, interactive=True), gr.Textbox(label="Disliked visual attributes", lines=1, value=negative_vp, interactive=True), gr.Button("Run personalized generation", scale=0, interactive=True, variant="primary")
 
@@ -539,8 +528,6 @@ def generate(prompt, vp_pos, vp_neg, slider, example_prompt, gallery, num_infere
             object_fit="contain", 
             height=500)
 
-
-
     return image, example_prompt, gallery
 
 def change_vp(extract_vp):
@@ -690,9 +677,6 @@ with gr.Blocks(css=css, title="ViPer Demo", theme=gr.themes.Base()) as demo:
             slider = gr.Slider(value=0.85, minimum=0, maximum=1.5, label="Personalization degree", interactive=True)
     
             with gr.Row():
-            
-                
-
                 prompt = gr.Dropdown(
                     example_prompts, label="Prompt", info="Enter your prompt or choose an example prompts from the dropdown.", allow_custom_value=True, multiselect=False, show_label=False
                 )
@@ -723,9 +707,7 @@ with gr.Blocks(css=css, title="ViPer Demo", theme=gr.themes.Base()) as demo:
                 randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
     with gr.Row(elem_id="main-container"):
-        with gr.Accordion("Images generated from the example prompts, but with different extracted preferences. The first image shows the non-personalized baseline generation.", open=False):
-            
-            
+        with gr.Accordion("Images generated from the example prompts, but with different extracted preferences. The first image shows the non-personalized baseline generation.", open=False): 
             example_prompt = gr.Markdown(f"Prompt: {example_prompts[0]}")
             gallery = gr.Gallery(
                 value=examples[example_prompts[0]],
@@ -740,7 +722,6 @@ with gr.Blocks(css=css, title="ViPer Demo", theme=gr.themes.Base()) as demo:
                 pre_prompt_button = gr.Button("⬅ Previous prompt", scale=1, interactive=True)
                 next_prompt_button = gr.Button("Next prompt ➡", scale=1, interactive=True)
             
-    
     submit_comment_button.click(
         fn = submit_comment,
         inputs = [comment],