Update app.py

app.py

@@ -16,7 +16,7 @@ import spaces
 from profanityfilter import ProfanityFilter
 
 import torch
-from diffusers import DiffusionPipeline, AutoencoderKL, DPMSolverMultistepScheduler
+from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
 
 # Set device
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -58,13 +58,116 @@ def prepare_image_for_watermark(image):
     background.save('custom_tabi.jpg')
 
     return background
+
+# Initialize the profanity filter
+pf = ProfanityFilter()
+def filter_inappropriate(input_text):
+    # Filter out inappropriate words
+    pf.censor_char = ' '
+    filtered_text = pf.censor(input_text)
+    return filtered_text.strip()
+
+# find the closest color name to rgb value
+def closest_color(rgb_color):
+    min_colors = {}
+    for key, name in webcolors.CSS3_HEX_TO_NAMES.items():
+        r_c, g_c, b_c = webcolors.hex_to_rgb(key)
+        rd = (r_c - rgb_color[0]) ** 2
+        gd = (g_c - rgb_color[1]) ** 2
+        bd = (b_c - rgb_color[2]) ** 2
+        min_colors[(rd + gd + bd)] = name
+    return min_colors[min(min_colors.keys())]
+
+def get_dominant_colors(img_filepath):
+  # Load the image from file path
+  img_data = Image.open(img_filepath)
+
+  # Convert the image to a NumPy array
+  img = np.array(img_data)
+
+  # k-means clustering to create palette of most dominant n_colors
+  pixels = np.float32(img.reshape(-1, 3))
+
+  n_colors = 2
+  criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, .1)
+  flags = cv2.KMEANS_RANDOM_CENTERS
+
+  _, labels, palette = cv2.kmeans(pixels, n_colors, None, criteria, 10, flags)
+
+  # get names of dominant colors
+  dominant_colors = []
+  for color in palette:
+        color_name = closest_color(color)
+        dominant_colors.append(color_name)
+
+  return dominant_colors
+
+def get_image_caption(image):
+    client = Client("https://vikhyatk-moondream1.hf.space/")
+    result = client.predict(
+		image,	# filepath  in 'image' Image component
+		"Describe the colors, patterns, aesthetic, artistic style, and objects in this photo",	# str  in 'Question' Textbox component
+		api_name="/answer_question"
+    )
+    print(result)
+    return result
+
+def get_image_keywords(image):
+  # get img2text description
+  caption = get_image_caption(image)
+
+  # get colors
+  colors_list = get_dominant_colors(image)
+  colors = ", ".join(colors_list)
+  return caption, colors
+
+from transformers import pipeline
+pipe = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-beta", torch_dtype=torch.bfloat16, device_map="auto")
+
+@spaces.GPU(enable_queue=True)
+def construct_prompt(image_caption, image_colors, user_input):
+    agent_maker_sys = f"""
+You are a AI whose job is to help users create their own custom shoe image which will reflect the characteristics, aesthetics, or include objects from an image described by users.
+In particular, you need to respond succintly and write a prompt for an image generation model. The response must include to the word "mm-tabi" which will trigger the style of shoe. The response should avoid any descriptions of man or woman and don't include any articles of clothing from Caption.
+The response should always end with "still life in the style of retrofuturism, highly detailed, hd, 8k". 
+The response should only use one or two aspects from the Caption provided by the user that could easily be applied to a still life scene or characteristic of the shoe, like color or texture or an object. 
+For example, if a user says, 
+"Colors: gainsboro, sienna
+/n Keywords: summer casual chic 
+/n Caption: The image features a man wearing a red and yellow sweater, which has a leopard print design. He is posing for the camera, and the sweater is placed on a mannequin, adding a touch of artistic style to the scene. The photo also includes a book, which is placed on a surface, possibly a table or a shelf. The book's presence adds an element of interest and contrast to the overall aesthetic of the image. The combination of the man in the sweater, the mannequin, and the book creates a visually appealing and artistic composition."
+, provide immediately an image prompt that describes a still life photo of a shoe corresponding to the keywords, color, and objects or stylistic elements from the caption provided.
+Immediately STOP after that. It should be in this format:
+"surreal photo of mm-tabi boot with split toe, surrounded by summer casual chic, red and yellow knit with leopard print, still life in the style of retrofuturism, fantasy, pimped, gainsboro, darkgray, darkolivegreen, sienna, highly detailed, hd, 8k"
+
+Here's another example, if a user says,
+"Colors: darkslategray, linen
+/n Keywords: baroque dystopia 
+/n Caption: The image features a young man wearing a jacket and a pair of black sunglasses. The jacket has a gray color with a patterned design, and the sunglasses are black and gray as well. The man is posing for the camera, and his facial expression is neutral. The overall aesthetic of the photo is casual and laid-back, with a focus on the sunglasses and the jacket as the main objects in the scene."
+; then respond:
+"surreal photo of mm-tabi boot with split toe, surrounded by a baroque dystopia, gray patterned design, casual and laid-back, still life in the style of retrofuturism, dreamy, unconventional, darkslategray, linen, silver, highly detailed, hd, 8k"
+"""
+
+    instruction = f"""
+<|system|>
+{agent_maker_sys}</s>
+<|user|>
+"""
+
+    prompt = f"{instruction.strip()}\n Colors: {image_colors} \n Keywords: {user_input} \n Caption: {image_caption}</s>"
+    print(f"PROMPT: {prompt}")
+    outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
+    
+    pattern = r'\<\|system\|\>(.*?)\<\|assistant\|\>'
+    cleaned_text = re.sub(pattern, '', outputs[0]["generated_text"], flags=re.DOTALL)
+
+    return cleaned_text.lstrip("\n")
     
 @spaces.GPU(enable_queue=True)
 def text2img_inference(prompt):
     gr.Info('Image generation request sent')
     with torch.no_grad():
         image = text2img_pipe(
-            prompt='surreal photo of a mm-tabi boot with split toe, still life in the style of retrofuturism, fantasy, unconventional, highly detailed, hd, 8k',
+            prompt=prompt,
             negative_prompt=negative_prompt,
             width=1024,
             height=1024,
@@ -72,19 +175,36 @@ def text2img_inference(prompt):
             guidance_scale=7.5
         ).images[0]
 
-    return image, "this is linda's ai generated caption"
+    return image
+
+def generate_image(user_input, input_image, image_prompt):
+    if image_prompt == "":
+      # generate keywords from image
+      gr.Info('Starting to generate caption for input image')
+      img_caption, img_colors = get_image_keywords(input_image)
+    
+      # filter user input
+      gr.Info('Processing keywords for inappropriate language')
+      user_input = filter_inappropriate(prompt)
+    
+      # construct prompt from image caption, image colors, and user input
+      gr.Info('Consrtucting prompt')
+      full_prompt = construct_prompt(img_caption, img_colors, user_input)
+      print(f"FULL PROMPT: {full_prompt}")
+
+    prompt = image_prompt if image_prompt != "" else full_prompt
 
-def generate_image(keywords, input_image, image_prompt):
-  image, prompt = text2img_inference(keywords)
-  watermarkable_image = prepare_image_for_watermark(image)
-  return watermarkable_image, prompt
+    # text2img generation with full prompt construction
+    image = text2img_inference(prompt)
+    watermarkable_image = prepare_image_for_watermark(image)
+    return watermarkable_image, prompt
 
 gradio_app = gr.Interface(
     fn=generate_image,
     inputs=[gr.Text(label="User Keywords"), gr.Image(label="Input Image", type='pil'), gr.Text(label="Generated Prompt")],
     outputs=[gr.Image(label="Image Generation"), gr.Text(label="Image Prompt")],
     title="Custom Tabi",
-    description="Enter a prompt with the trigger word 'TOK' to indicate in the style of maison margiela tabi shoe. Ex. oil painting still life of TOK boot on the beach, burgers and fries"
+    description="Enter keywords and upload image to generate a custom Tabi boot"
 )
 
 gradio_app.launch(debug=True)