app.py

import gradio as gr
import numpy as np
import random
from diffusers import DiffusionPipeline
import torch
from diffusers import StableDiffusionXLPipeline
import requests
import torch
from PIL import Image
from transformers import AutoProcessor, AutoModelForVision2Seq, BitsAndBytesConfig
from transformers.image_utils import load_image
from peft import PeftModel
import re
from diffusers import StableDiffusionXLPipeline, DiffusionPipeline

device = "cuda" if torch.cuda.is_available() else "cpu"
print(device)

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    llm_int8_skip_modules=["lm_head", "embed_tokens"],
)

processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b", size= {"longest_edge": 448, "shortest_edge": 378}, do_image_splitting=False)

if torch.cuda.is_available():
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0", 
        torch_dtype=torch.float16
    ).to("cuda")
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

valid_api = ""

css="""
#col-container {
    margin: 0 auto;
    max-width: 520px;
}

#gen-container {
    margin: 0 auto;
    max-width: 640px;
}

#title-container {
    margin: 0 auto;
    max-width: 1340px;
}

#main-container {
    margin: 0 auto;
    max-width: 1340px;
}
"""


if torch.cuda.is_available():
    power_device = "GPU"
else:
    power_device = "CPU"

from PIL import Image 
  
comment_images = [
    "test.png",
    "comment_images/0.png",
    "comment_images/1.png",
    "comment_images/2.png",
    "comment_images/3.png",
    "comment_images/4.png",
    "comment_images/5.png",
    "comment_images/6.png",
    "comment_images/7.png",
    "comment_images/8.png",
    "comment_images/9.png",
    "comment_images/10.png",
    "comment_images/11.png",
    "comment_images/12.png",
    "comment_images/13.png",
    "comment_images/14.png",
    "comment_images/15.png",
    "comment_images/16.png",
    "comment_images/17.png",
    "comment_images/18.png",
    "comment_images/19.png",
    "comment_images/20.png",
    "comment_images/21.png",
    "comment_images/22.png",
    "comment_images/23.png",
    "comment_images/24.png",
    "comment_images/25.png",
    "comment_images/26.png",
    "comment_images/27.png",
    "comment_images/28.png",
    "comment_images/29.png",
    "comment_images/30.png",
    "comment_images/31.png",
    "comment_images/32.png",
    "comment_images/33.png",
    "comment_images/34.png",
    "comment_images/35.png",
    "comment_images/36.png",
    "comment_images/37.png",
    "comment_images/38.png",
    "comment_images/39.png",
    "comment_images/40.png",
    "comment_images/41.png",
    "comment_images/42.jpg",
    "comment_images/43.png",
    "comment_images/44.png",
    "comment_images/45.png",
    "comment_images/46.png",
    "comment_images/47.png",
    "comment_images/48.png",
    "comment_images/49.png",
    "comment_images/50.png",
    "comment_images/51.png",
    "comment_images/52.png",
    "comment_images/53.png",
    "comment_images/54.png",
    "comment_images/55.png"
]

comments = {'test.png': "Not sure about the concept, it's too straightforward. Though the boy looks kinda creepy which makes it exciting. the art style is pretty to look at. I like that the colors are muted, but wish they were a bit darker to make it more eerie and add depth.", 'comment_images/0.png': "Hate this with a passion. The colors are too vibrant and don't match at all. I hate these colors in general. The patterns are too abstract and contemporary. a 5-year-old could draw this. pass.", 'comment_images/1.png': "Woah I love the art style. The texture feels like old paper which is oh so beautiful. There are so many details to focus on. I love the expressive lines and how busy the composition is. Even though orange isn't my favorite, the greenish blue color of the water is so gorgeous.", 'comment_images/2.png': "I don't like how monochromatic and muted this one is. but the paperish texture is nice and the details are so intricate.", 'comment_images/3.png': "Oh super pretty! Looks so smooth and wet. Love the details and loose lines too. Feels mystical and magical and eerie. Also dark purples and blues? deep indigo? My fav ever. I'm here for it.", 'comment_images/4.png': "Love the art style. The uncanny vibe and nightmarish horror is so cool. Like its horror but if you squint you can't tell? Love the strange. wish it had more colors though. not a fan of greyscale.", 'comment_images/5.png': 'omg I hate this haha. what the hell. everything about it disgusts me so boring and childish ew.', 'comment_images/6.png': 'yessss. give it to the texture give it to the brushstrokes give it to the style. perfect. just wish the colors were less beige and more bold. I want an active nightmare. but kisses to the surrealism.'}
comments = dict()

image_index = 0

def submit_comment(comment):
    global comment_images, image_index
    if 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:
        comments.pop(comment_images[0], None)

    print(comments)
    next_comment = ""
    if comment_images[0] in comments:
        next_comment = comments[comment_images[0]]

    clear_botton = gr.Button("Clear comments", interactive=len(comments) != 0)
        
    return (gr.Image(value=comment_images[0], label=f"image {image_index+1}/{len(comment_images)}", show_label=True),
            gr.Text(label="Comment", show_label=False, lines=2, max_lines=3, placeholder="Enter your comment", value=next_comment, container=False),
            gr.Button(f"Extract visual preference from {len(comments)} comments", interactive=len(comments) != 0),
            clear_botton
           )

def next_image():
    global comment_images, image_index
    comment_images.append(comment_images[0])
    comment_images = comment_images[1:]

    print(comments)
    next_comment = ""
    if comment_images[0] in comments:
        next_comment = comments[comment_images[0]]

    image_index = (image_index + 1) % len(comment_images)
        
    return gr.Image(value=comment_images[0], label=f"image {image_index+1}/{len(comment_images)}", show_label=True), gr.Text(label="Comment", show_label=False, lines=2, max_lines=3, placeholder="Enter your comment", value=next_comment, container=False)

def previous_image():
    global comment_images, image_index
    comment_images = comment_images[::-1]
    comment_images.append(comment_images[0])
    comment_images = comment_images[1:]
    comment_images = comment_images[::-1]

    print(comments)
    next_comment = ""
    if comment_images[0] in comments:
        next_comment = comments[comment_images[0]]

    image_index = (image_index - 1) % len(comment_images)
        
    return gr.Image(value=comment_images[0], label=f"image {image_index+1}/{len(comment_images)}", show_label=True), gr.Text(label="Comment", show_label=False, lines=2, max_lines=3, placeholder="Enter your comment", value=next_comment, container=False)

def clear_comments():
    comments.clear()
    extract_vp_botton = gr.Button(f"Extract visual preference from {len(comments)} comments", interactive=len(comments) != 0)
    clear_botton = gr.Button("Clear comments", interactive=len(comments) != 0)
    return extract_vp_botton, clear_botton
    

def extract_vp():
    if valid_api == "":
        vpe_model = AutoModelForVision2Seq.from_pretrained(
            "HuggingFaceM4/idefics2-8b",
            torch_dtype=torch.float16,    
            quantization_config=bnb_config,
        )

        vpe_model = PeftModel.from_pretrained(vpe_model, "VPE2").to("cuda")
    
        global comments
        
        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.

If the person expresses a preference for a specific aspect without clearly stating its category (e.g., appreciating the colors without specifying which colors), identify these specific features from the images directly to make the person's preference understandable without needing to see the artwork.

Your output should consist of two concise lists of keywords: one listing the specific art features the person likes and another listing the specific features they dislike (specified in keyword format without using sentences).

Here are the images and their corresponding comments:
"""
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", 
                    "text": prompt},
                ]
            }
        ]
        images = []
        comment_number = 1
        for image in comments:
            comment = comments[image]
            image = Image.open(image)
            images.append(image)
    
            messages[0]["content"].append(
                {"type": "image"}
            )

            messages[0]["content"].append(
                {"type": "text", 
                 "text": f"Comment {comment_number}: {comment}"}
            )
            comment_number = comment_number + 1
    
        prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
        inputs = processor(text=prompt, images=images, return_tensors="pt")
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        generated_ids = vpe_model.generate(**inputs, max_new_tokens=2000, repetition_penalty=0.99, do_sample=False)
        del vpe_model
        generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
        positive_vp, negative_vp = re.search('.* \nAssistant: Liked Art Features: (.*)\nDisliked Art Features: (.*)', generated_texts).groups()

    else:
        PRINT(valid_api)
    
    gr.Info("Visual preference successfully extracted.")
    
    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", scale=0, interactive=True)

def api_fn(api):
    global valid_api
    if api != "correct":
        gr.Warning("Invalid API!")
        valid_api = ""
    else:
        gr.Info("Valid API")
        valid_api = api

def generate(prompt, vp_pos, vp_neg, slider):
    print(f"prompt: {prompt}")
    image = pipe(prompt=prompt, 
                num_inference_steps=40, 
                vp_pos=vp_pos, 
                vp_neg=vp_neg, 
                vp_degree_pos=slider,
                vp_degree_neg=slider
    ).images[0]
    return image

def change_vp(extract_vp):
    return

def upload_file(files):
    global comment_images, image_index
    file_path = [file.name for file in files][0]
    comment_images = [file_path] + comment_images

    next_comment = ""
    return gr.Image(value=comment_images[0], label=f"image {image_index+1}/{len(comment_images)}", show_label=True), gr.Text(label="Comment", show_label=False, lines=2, max_lines=3, placeholder="Enter your comment", value=next_comment, container=False)

with gr.Blocks(css=css, title="ViPer Demo", theme=gr.themes.Base()) as demo:
    with gr.Row(elem_id="title-container"):
        gr.Markdown(f"""
                # **ViPer: Visual Personalization of Generative Models via Individual Preference Learning**
                \n
                \n
                \n

                """)
    with gr.Row(elem_id="main-container"):
        with gr.Column(elem_id="col-container"):
    
            gr.Markdown(f"""
                ## Step 1: Extracting visual preference from comments on images
            """
            )

            gr.Markdown("Please write your comments on the images below, explaining why you like or dislike each one from an artistic perspective. Focus on images that evoke **strong reactions**, whether positive or negative, and skip those that don't affect you much.\nMore **detailed** comments will help us provide more personalized results. We recommend commenting on **at least 8** images.")

            gr.Markdown("Note that our method works best with an OpenAI API. The free method might result in minor hallucinations in the extracted visual preferences.")

            with gr.Accordion("Examples of Effective Comments", open=False):
                example_comment_1 = gr.Textbox(
                    label="Example 1",
                    lines=4,
                    value="Gotta say I love this one. The idea of collage painting really appeals to me. I can pick up on the subtle shadows. The combination of soft, creamy yellow and warm green looks really nice too. The paper texture itself is really interesting.",
                )

                example_comment_2 = gr.Textbox(
                    label="Example 2",
                    lines=4,
                    value="I adore the blue and greenish-blue palette, blue Dianne, and dark colors of this image. I also appreciate the Hergé inspiration in this artwork. However, I would have preferred a more complex and adventurous concept rather than a simple landscape. I wish it was more surreal and creepy.",
                )
    
            comment_image = gr.Image(value=comment_images[0], label=f"image {image_index+1}/{len(comment_images)}", show_label=True)
                      
            comment = gr.Text(
                label="Comment",
                show_label=False,
                lines=2,
                max_lines=3,
                placeholder="Enter your comment",
                container=False,
            )
    
            with gr.Row():
                submit_comment_button = gr.Button("Submit comment", scale=0)
                previous_image_botton = gr.Button("Previous Image", scale=0)
                next_image_botton = gr.Button("Next Image", scale=0)

            file_output = gr.File(visible=False)
            upload_button = gr.UploadButton("Click to upload images", file_types=["image"], file_count="multiple")
            
            clear_botton = gr.Button("Clear comments", interactive=len(comments) != 0)

            with gr.Accordion("Enter GPT API for Better Results (optional)", open=False):
                with gr.Row():
                    api = gr.Text(
                        max_lines=1,
                        placeholder="Enter your API",
                        container=False,
                    )
                    
                    api_button = gr.Button("Enter", scale=0)
    
            extract_vp_botton = gr.Button(f"Extract visual preference from {len(comments)} comments", interactive=len(comments) != 0)

        with gr.Column(elem_id="gen-container"):
            gr.Markdown(f"""
                You can edit your visual preference in case of hallucinations.
            """
            )
            
            positive_extracted_vp = gr.Textbox(
                label="Liked visual attributes",
                lines=3,
                value="",
            )

            negative_extracted_vp = gr.Textbox(
                label="Disliked visual attributes",
                lines=1,
                value="",
            )
    
            gr.Markdown(f"""
                ## Step 2: Personalized image generation (using Stable Diffusion XL)
                Write down the prompt to generate your preferred images once your visual preference has been extracted from your comments.
                """)

            slider = gr.Slider(value=0.85, minimum=0, maximum=1.5, label="Personalization degree", interactive=True)
    
            with gr.Row():
            
                prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter your prompt",
                    container=False,
                )
                
                run_button = gr.Button("Run", scale=0, interactive=False)

            result = gr.Image(label="Result", show_label=False, interactive=False, value="Train running through the coun (6).png")

    with gr.Row(elem_id="main-container"):
        with gr.Accordion("images generated from the same prompt but different extracted preferences (prompt: Picture of a lady)", open=False):
            examples = [
                "examples/(0).png",
                "examples/(16).png",
                "examples/(2).png",
                "examples/(12).png",
                "examples/(13).png",
                "examples/(14).png",
                "examples/(15).png",
                "examples/(17).png",
                "examples/(11).png",
                "examples/(18).png",
            ]
            gallery = gr.Gallery(
                value=examples,
                label="", 
                show_label=False,
                columns=[5],
                rows=[2], 
                object_fit="contain", 
                height=500)
            
    
    submit_comment_button.click(
        fn = submit_comment,
        inputs = [comment],
        outputs = [comment_image, comment, extract_vp_botton, clear_botton]
    )

    previous_image_botton.click(
        fn = previous_image,
        inputs = [],
        outputs = [comment_image, comment]
    )

    next_image_botton.click( 
        fn = next_image,
        inputs = [],
        outputs = [comment_image, comment]
    )

    extract_vp_botton.click(
        fn = extract_vp,
        inputs = [],
        outputs = [positive_extracted_vp, negative_extracted_vp, run_button]
    )

    api_button.click(
        fn = api_fn,
        inputs = [api],
        outputs = [],
    )

    run_button.click(
        fn = generate,
        inputs = [prompt, positive_extracted_vp, negative_extracted_vp, slider],
        outputs = [result],
    )

    positive_extracted_vp.change(
        fn = change_vp,
        inputs = [positive_extracted_vp],
        outputs = [],
    )
    
    negative_extracted_vp.change(
        fn = change_vp,
        inputs = [negative_extracted_vp],
        outputs = [],
    )

    clear_botton.click(
        fn = clear_comments,
        inputs = [],
        outputs = [extract_vp_botton, clear_botton]
    )

    upload_button.upload(
        upload_file, 
        upload_button, 
        [comment_image, comment]
    )

demo.launch(share=True)