Spaces:
Runtime error
Runtime error
Update caption_generator.py
Browse files- caption_generator.py +48 -0
caption_generator.py
CHANGED
|
@@ -36,6 +36,54 @@ EXTRA_OPTIONS = [
|
|
| 36 |
"ONLY describe the most important elements of the image."
|
| 37 |
]
|
| 38 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 39 |
# Image Adapter
|
| 40 |
class ImageAdapter(nn.Module):
|
| 41 |
def __init__(self, input_features: int, output_features: int, ln1: bool, pos_emb: bool, num_image_tokens: int, deep_extract: bool):
|
|
|
|
| 36 |
"ONLY describe the most important elements of the image."
|
| 37 |
]
|
| 38 |
|
| 39 |
+
CAPTION_TYPE_MAP = {
|
| 40 |
+
"Descriptive": [
|
| 41 |
+
"Write a descriptive caption for this image in a formal tone.",
|
| 42 |
+
"Write a descriptive caption for this image in a formal tone within {word_count} words.",
|
| 43 |
+
"Write a {length} descriptive caption for this image in a formal tone.",
|
| 44 |
+
],
|
| 45 |
+
"Descriptive (Informal)": [
|
| 46 |
+
"Write a descriptive caption for this image in a casual tone.",
|
| 47 |
+
"Write a descriptive caption for this image in a casual tone within {word_count} words.",
|
| 48 |
+
"Write a {length} descriptive caption for this image in a casual tone.",
|
| 49 |
+
],
|
| 50 |
+
"Training Prompt": [
|
| 51 |
+
"Write a stable diffusion prompt for this image.",
|
| 52 |
+
"Write a stable diffusion prompt for this image within {word_count} words.",
|
| 53 |
+
"Write a {length} stable diffusion prompt for this image.",
|
| 54 |
+
],
|
| 55 |
+
"MidJourney": [
|
| 56 |
+
"Write a MidJourney prompt for this image.",
|
| 57 |
+
"Write a MidJourney prompt for this image within {word_count} words.",
|
| 58 |
+
"Write a {length} MidJourney prompt for this image.",
|
| 59 |
+
],
|
| 60 |
+
"Booru tag list": [
|
| 61 |
+
"Write a list of Booru tags for this image.",
|
| 62 |
+
"Write a list of Booru tags for this image within {word_count} words.",
|
| 63 |
+
"Write a {length} list of Booru tags for this image.",
|
| 64 |
+
],
|
| 65 |
+
"Booru-like tag list": [
|
| 66 |
+
"Write a list of Booru-like tags for this image.",
|
| 67 |
+
"Write a list of Booru-like tags for this image within {word_count} words.",
|
| 68 |
+
"Write a {length} list of Booru-like tags for this image.",
|
| 69 |
+
],
|
| 70 |
+
"Art Critic": [
|
| 71 |
+
"Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc.",
|
| 72 |
+
"Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc. Keep it within {word_count} words.",
|
| 73 |
+
"Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc. Keep it {length}.",
|
| 74 |
+
],
|
| 75 |
+
"Product Listing": [
|
| 76 |
+
"Write a caption for this image as though it were a product listing.",
|
| 77 |
+
"Write a caption for this image as though it were a product listing. Keep it under {word_count} words.",
|
| 78 |
+
"Write a {length} caption for this image as though it were a product listing.",
|
| 79 |
+
],
|
| 80 |
+
"Social Media Post": [
|
| 81 |
+
"Write a caption for this image as if it were being used for a social media post.",
|
| 82 |
+
"Write a caption for this image as if it were being used for a social media post. Limit the caption to {word_count} words.",
|
| 83 |
+
"Write a {length} caption for this image as if it were being used for a social media post.",
|
| 84 |
+
],
|
| 85 |
+
}
|
| 86 |
+
|
| 87 |
# Image Adapter
|
| 88 |
class ImageAdapter(nn.Module):
|
| 89 |
def __init__(self, input_features: int, output_features: int, ln1: bool, pos_emb: bool, num_image_tokens: int, deep_extract: bool):
|