pratikshahp commited on
Commit
9a101a6
Β·
verified Β·
1 Parent(s): ec03bd7

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +44 -0
app.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import pipeline
4
+
5
+ # Load the models
6
+ caption_generator = pipeline("image-captioning", model="gokaygokay/SD3-Long-Captioner")
7
+ compliment_generator = pipeline("text-generation", model="hysts/zephyr-7b")
8
+
9
+ SYSTEM_PROMPT = """
10
+ You are helpful assistant that gives the best compliments to people.
11
+ You will be given a caption of someone's headshot.
12
+ Based on that caption, provide a one sentence compliment to the person in the image.
13
+ Make sure you compliment the person in the image and not any objects or scenery.
14
+ Do NOT include any hashtags in your compliment or phrases like (emojis: dog, smiling face with heart-eyes, sun).
15
+ Here are some examples of the desired behavior:
16
+ Caption: a front view of a man who is smiling, there is a lighthouse in the background, there is a grassy area on the left that is green and curved. in the distance you can see the ocean and the shore. there is a grey and cloudy sky above the lighthouse and the trees.
17
+ Compliment: Your smile is as bright as a lighthouse, lighting up the world around you. 🌟
18
+ Caption: in a close-up, a blonde woman with short, wavy hair, is the focal point of the image. she's dressed in a dark brown turtleneck sweater, paired with a black hat and a black suit jacket. her lips are a vibrant red, and her eyes are a deep brown. in the background, a man with a black hat and a white shirt is visible.
19
+ Compliment: You are the epitome of elegance and grace, with a style that is as timeless as your beauty. πŸ’ƒπŸŽ©
20
+ Conversation begins below:
21
+ """
22
+
23
+ def generate_compliment(image):
24
+ # Generate caption
25
+ caption = caption_generator(image)[0]['caption']
26
+ # Generate compliment
27
+ input_text = f"Caption: {caption}\nCompliment: "
28
+ generated_compliment = compliment_generator(SYSTEM_PROMPT + input_text, max_new_tokens=50, do_sample=True)[0]['generated_text']
29
+ compliment = generated_compliment.split('Compliment: ')[-1].strip()
30
+ return caption, compliment
31
+
32
+ # Gradio interface
33
+ iface = gr.Interface(
34
+ fn=generate_compliment,
35
+ inputs=gr.inputs.Image(type="pil"),
36
+ outputs=[
37
+ gr.outputs.Textbox(label="Caption"),
38
+ gr.outputs.Textbox(label="Compliment")
39
+ ],
40
+ title="Compliment Bot πŸ’–",
41
+ description="Upload your headshot and get a personalized compliment!"
42
+ )
43
+
44
+ iface.launch()