app.py

import spaces
import gradio as gr
from huggingface_hub import InferenceClient
from torch import nn
from transformers import AutoModel, AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM
from pathlib import Path
import torch
import torch.amp.autocast_mode
from PIL import Image
import os
import torchvision.transforms.functional as TVF

from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

USERNAME = os.getenv("USERNAME")
PASSWORD = os.getenv("PASSWORD")

CLIP_PATH = "google/siglip-so400m-patch14-384"
MODEL_PATH = "meta-llama/Meta-Llama-3.1-8B"
CHECKPOINT_PATH = Path("cgrkzexw-599808")
TITLE = "<h1><center>JoyCaption Alpha One (2024-09-20a)</center></h1>"
CAPTION_TYPE_MAP = {
    "Descriptive": [
        "Write a descriptive caption for this image in a formal tone.",
        "Write a descriptive caption for this image in a formal tone within {word_count} words.",
        "Write a {length} descriptive caption for this image in a formal tone.",
    ],
    "Descriptive (Informal)": [
        "Write a descriptive caption for this image in a casual tone.",
        "Write a descriptive caption for this image in a casual tone within {word_count} words.",
        "Write a {length} descriptive caption for this image in a casual tone.",
    ],
    "Stable Diffusion Prompt": [
        "Write a stable diffusion prompt for this image.",
        "Write a stable diffusion prompt for this image within {word_count} words.",
        "Write a {length} stable diffusion prompt for this image.",
    ],
    "MidJourney": [
        "Write a MidJourney prompt for this image.",
        "Write a MidJourney prompt for this image within {word_count} words.",
        "Write a {length} MidJourney prompt for this image.",
    ],
    "Booru tag list": [
        "Write a list of Booru tags for this image.",
        "Write a list of Booru tags for this image within {word_count} words.",
        "Write a {length} list of Booru tags for this image.",
    ],
    "Booru-like tag list": [
        "Write a list of Booru-like tags for this image.",
        "Write a list of Booru-like tags for this image within {word_count} words.",
        "Write a {length} list of Booru-like tags for this image.",
    ],
    "Art Critic": [
        "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.",
        "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.",
        "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}.",
    ],
    "Product Listing": [
        "Write a caption for this image as though it were a product listing.",
        "Write a caption for this image as though it were a product listing. Keep it under {word_count} words.",
        "Write a {length} caption for this image as though it were a product listing.",
    ],
    "Social Media Post": [
        "Write a caption for this image as if it were being used for a social media post.",
        "Write a caption for this image as if it were being used for a social media post. Limit the caption to {word_count} words.",
        "Write a {length} caption for this image as if it were being used for a social media post.",
    ],

    "Style Prompt": [
        "Generate a detailed style prompt for this image incorporating the specified lens type, film stock, composition style, lighting aspect, special technique, and color effect.",
        "Generate a detailed style prompt for this image incorporating the specified lens type, film stock, composition style, lighting aspect, special technique, and color effect. Limit the prompt to {word_count} words.",
        "Generate a {length} detailed style prompt for this image incorporating the specified lens type, film stock, composition style, lighting aspect, special technique, and color effect.",
    ],
}

HF_TOKEN = os.environ.get("HF_TOKEN", None)


class ImageAdapter(nn.Module):
	def __init__(self, input_features: int, output_features: int, ln1: bool, pos_emb: bool, num_image_tokens: int, deep_extract: bool):
		super().__init__()
		self.deep_extract = deep_extract

		if self.deep_extract:
			input_features = input_features * 5

		self.linear1 = nn.Linear(input_features, output_features)
		self.activation = nn.GELU()
		self.linear2 = nn.Linear(output_features, output_features)
		self.ln1 = nn.Identity() if not ln1 else nn.LayerNorm(input_features)
		self.pos_emb = None if not pos_emb else nn.Parameter(torch.zeros(num_image_tokens, input_features))

		# Mode token
		#self.mode_token = nn.Embedding(n_modes, output_features)
		#self.mode_token.weight.data.normal_(mean=0.0, std=0.02)   # Matches HF's implementation of llama3

		# Other tokens (<|image_start|>, <|image_end|>, <|eot_id|>)
		self.other_tokens = nn.Embedding(3, output_features)
		self.other_tokens.weight.data.normal_(mean=0.0, std=0.02)   # Matches HF's implementation of llama3

	def forward(self, vision_outputs: torch.Tensor):
		if self.deep_extract:
			x = torch.concat((
				vision_outputs[-2],
				vision_outputs[3],
				vision_outputs[7],
				vision_outputs[13],
				vision_outputs[20],
			), dim=-1)
			assert len(x.shape) == 3, f"Expected 3, got {len(x.shape)}"  # batch, tokens, features
			assert x.shape[-1] == vision_outputs[-2].shape[-1] * 5, f"Expected {vision_outputs[-2].shape[-1] * 5}, got {x.shape[-1]}"
		else:
			x = vision_outputs[-2]

		x = self.ln1(x)

		if self.pos_emb is not None:
			assert x.shape[-2:] == self.pos_emb.shape, f"Expected {self.pos_emb.shape}, got {x.shape[-2:]}"
			x = x + self.pos_emb

		x = self.linear1(x)
		x = self.activation(x)
		x = self.linear2(x)

		# Mode token
		#mode_token = self.mode_token(mode)
		#assert mode_token.shape == (x.shape[0], mode_token.shape[1], x.shape[2]), f"Expected {(x.shape[0], 1, x.shape[2])}, got {mode_token.shape}"
		#x = torch.cat((x, mode_token), dim=1)

		# <|image_start|>, IMAGE, <|image_end|>
		other_tokens = self.other_tokens(torch.tensor([0, 1], device=self.other_tokens.weight.device).expand(x.shape[0], -1))
		assert other_tokens.shape == (x.shape[0], 2, x.shape[2]), f"Expected {(x.shape[0], 2, x.shape[2])}, got {other_tokens.shape}"
		x = torch.cat((other_tokens[:, 0:1], x, other_tokens[:, 1:2]), dim=1)

		return x

	def get_eot_embedding(self):
		return self.other_tokens(torch.tensor([2], device=self.other_tokens.weight.device)).squeeze(0)



# Load CLIP
print("Loading CLIP")
clip_processor = AutoProcessor.from_pretrained(CLIP_PATH)
clip_model = AutoModel.from_pretrained(CLIP_PATH)
clip_model = clip_model.vision_model

if (CHECKPOINT_PATH / "clip_model.pt").exists():
	print("Loading VLM's custom vision model")
	checkpoint = torch.load(CHECKPOINT_PATH / "clip_model.pt", map_location='cpu')
	checkpoint = {k.replace("_orig_mod.module.", ""): v for k, v in checkpoint.items()}
	clip_model.load_state_dict(checkpoint)
	del checkpoint

clip_model.eval()
clip_model.requires_grad_(False)
clip_model.to("cuda")


# Tokenizer
print("Loading tokenizer")
tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_PATH / "text_model", use_fast=True)
assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance(tokenizer, PreTrainedTokenizerFast), f"Tokenizer is of type {type(tokenizer)}"

# LLM
print("Loading LLM")
print("Loading VLM's custom text model")
text_model = AutoModelForCausalLM.from_pretrained(CHECKPOINT_PATH / "text_model", device_map=0, torch_dtype=torch.bfloat16)
text_model.eval()

# Image Adapter
print("Loading image adapter")
image_adapter = ImageAdapter(clip_model.config.hidden_size, text_model.config.hidden_size, False, False, 38, False)
image_adapter.load_state_dict(torch.load(CHECKPOINT_PATH / "image_adapter.pt", map_location="cpu"))
image_adapter.eval()
image_adapter.to("cuda")


def preprocess_image(input_image: Image.Image) -> torch.Tensor:
    image = input_image.resize((384, 384), Image.LANCZOS)
    pixel_values = TVF.pil_to_tensor(image).unsqueeze(0) / 255.0
    pixel_values = TVF.normalize(pixel_values, [0.5], [0.5])
    return pixel_values.to('cuda')

def generate_caption(text_model, tokenizer, image_features, prompt_str: str, max_new_tokens: int = 300) -> str:
    prompt = tokenizer.encode(prompt_str, return_tensors='pt', padding=False, truncation=False, add_special_tokens=False)
    prompt_embeds = text_model.model.embed_tokens(prompt.to('cuda'))

    convo = [
        {"role": "system", "content": "You are a helpful image captioner."},
        {"role": "user", "content": prompt_str},
    ]
    convo_string = tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
    convo_tokens = tokenizer.encode(convo_string, return_tensors="pt", add_special_tokens=False, truncation=False)
    convo_tokens = convo_tokens.squeeze(0)

    eot_id_indices = (convo_tokens == tokenizer.convert_tokens_to_ids("<|eot_id|>")).nonzero(as_tuple=True)[0].tolist()
    assert len(eot_id_indices) == 2, f"Expected 2 <|eot_id|> tokens, got {len(eot_id_indices)}"
    preamble_len = eot_id_indices[1] - prompt.shape[1]

    convo_embeds = text_model.model.embed_tokens(convo_tokens.unsqueeze(0).to('cuda'))

    input_embeds = torch.cat([
        convo_embeds[:, :preamble_len],
        image_features.to(dtype=convo_embeds.dtype),
        convo_embeds[:, preamble_len:],
    ], dim=1).to('cuda')

    input_ids = torch.cat([
        convo_tokens[:preamble_len].unsqueeze(0),
        torch.zeros((1, image_features.shape[1]), dtype=torch.long),
        convo_tokens[preamble_len:].unsqueeze(0),
    ], dim=1).to('cuda')
    attention_mask = torch.ones_like(input_ids)

    generate_ids = text_model.generate(input_ids, inputs_embeds=input_embeds, attention_mask=attention_mask, max_new_tokens=max_new_tokens, do_sample=True, suppress_tokens=None)

    generate_ids = generate_ids[:, input_ids.shape[1]:]
    if generate_ids[0][-1] == tokenizer.eos_token_id or generate_ids[0][-1] == tokenizer.convert_tokens_to_ids("<|eot_id|>"):
        generate_ids = generate_ids[:, :-1]

    return tokenizer.batch_decode(generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)[0]

@spaces.GPU()
@torch.no_grad()
def stream_chat(input_image: Image.Image, caption_type: str, caption_length: str | int, extra_options: list[str], name_input: str, custom_prompt: str, lens_type: str = "", film_stock: str = "", composition_style: str = "", lighting_aspect: str = "", special_technique: str = "", color_effect: str = "") -> tuple[str, str]:
    torch.cuda.empty_cache()

    # 'any' means no length specified
    length = None if caption_length == "any" else caption_length

    if isinstance(length, str):
        try:
            length = int(length)
        except ValueError:
            pass
    
    # Build prompt
    if length is None:
        map_idx = 0
    elif isinstance(length, int):
        map_idx = 1
    elif isinstance(length, str):
        map_idx = 2
    else:
        raise ValueError(f"Invalid caption length: {length}")
    
    prompt_str = CAPTION_TYPE_MAP[caption_type][map_idx]

    # Add extra options
    if len(extra_options) > 0:
        prompt_str += " " + " ".join(extra_options)
    
    # Add name, length, word_count
    prompt_str = prompt_str.format(name=name_input, length=caption_length, word_count=caption_length)

    if custom_prompt.strip() != "":
        prompt_str = custom_prompt.strip()
    
    # Add style prompt options
    if caption_type == "Style Prompt":
        prompt_str += f" Incorporate the effect of a {lens_type} lens. "
        prompt_str += f"Apply the characteristics of {film_stock} film stock. "
        prompt_str += f"Use a {composition_style} composition style. "
        prompt_str += f"Implement {lighting_aspect} lighting. "
        prompt_str += f"Apply the {special_technique} technique. "
        prompt_str += f"Use a {color_effect} color effect. "
    
    # For debugging
    print(f"Prompt: {prompt_str}")

    # Preprocess image
    pixel_values = preprocess_image(input_image)

    # Embed image
    with torch.amp.autocast_mode.autocast('cuda', enabled=True):
        vision_outputs = clip_model(pixel_values=pixel_values, output_hidden_states=True)
        image_features = vision_outputs.hidden_states
        embedded_images = image_adapter(image_features)
        embedded_images = embedded_images.to('cuda')

    # Build the conversation
    convo = [
        {
            "role": "system",
            "content": "You are a helpful image captioner.",
        },
        {
            "role": "user",
            "content": prompt_str,
        },
    ]

    # Format the conversation
    convo_string = tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
    assert isinstance(convo_string, str)

    # Generate caption
    caption = generate_caption(text_model, tokenizer, embedded_images, prompt_str)

    return prompt_str, caption.strip()

css = """
h1, h2, h3, h4, h5, h6, p, li, ul, ol, a, img {
    text-align: left;
}
img {
    display: inline-block;
    vertical-align: middle;
    margin-right: 10px;
    max-width: 100%;
    height: auto;
}
.centered-image {
    display: block;
    margin-left: auto;
    margin-right: auto;
    max-width: 100%;
    height: auto;
}
ul, ol {
    padding-left: 20px;
}
.gradio-container {
    max-width: 100% !important;
    padding: 0 !important;
}
.gradio-row {
    margin-left: 0 !important;
    margin-right: 0 !important;
}
.gradio-column {
    padding-left: 0 !important;
    padding-right: 0 !important;
}
/* Left-align dropdown text */
.gradio-dropdown > div {
    text-align: left !important;
}
/* Left-align checkbox labels */
.gradio-checkbox label {
    text-align: left !important;
}
/* Left-align radio button labels */
.gradio-radio label {
    text-align: left !important;
}
/* Change button color to green */
.gradio-button {
    background-color: #4CAF50 !important; /* Green color */
    border-color: #4CAF50 !important;
    color: white !important;
}
.gradio-button:hover {
    background-color: #45a049 !important; /* Darker green on hover */
    border-color: #45a049 !important;
}
"""

# Add detailed descriptions for each option
lens_types_info = {
    "Standard": "A versatile lens with a field of view similar to human vision.",
    "Wide-angle": "Captures a wider field of view, great for landscapes and architecture. Applies moderate to strong lens effect with image warp.",
    "Telephoto": "Used for distant subjects, gives an 'award-winning' or 'National Geographic' look. Creates interesting effects when prompted.",
    "Macro": "For extreme close-up photography, revealing tiny details.",
    "Fish-eye": "Ultra-wide-angle lens that creates a strong bubble-like distortion. Generates panoramic photos with the entire image warping into a bubble.",
    "Tilt-shift": "Allows adjusting the plane of focus, creating a 'miniature' effect. Known for the 'diorama miniature look'.",
    "Zoom lens": "Variable focal length lens. Often zooms in on the subject, perfect for creating a base for inpainting. Interesting effect on landscapes with motion blur.",
    "GoPro": "Wide-angle lens with clean digital look. Excludes film grain and most filter effects, resulting in natural colors and regular saturation.",
    "Pinhole camera": "Creates a unique, foggy, low-detail, historic photograph look. Used since the 1850s, with peak popularity in the 1930s."
}

film_stocks_info = {
    "Kodak Portra": "Professional color negative film known for its natural skin tones and low contrast.",
    "Fujifilm Velvia": "Slide film known for vibrant colors and high saturation, popular among landscape photographers.",
    "Ilford Delta": "Black and white film known for its fine grain and high sharpness.",
    "Kodak Tri-X": "Classic high-speed black and white film, known for its distinctive grain and wide exposure latitude.",
    "Fujifilm Provia": "Color reversal film known for its natural color reproduction and fine grain.",
    "Cinestill": "Color photos with fine/low grain and higher than average resolution. Colors are slightly oversaturated or slightly desaturated.",
    "Ektachrome": "Color photos with fine/low to moderate grain. Colors on the colder part of spectrum or regular, with normal or slightly higher saturation.",
    "Ektar": "Modern Kodak film. Color photos with little to no grain. Results look like regular modern photography with artistic angles.",
    "Film Washi": "Mostly black and white photos with fine/low to moderate grain. Occasionally gives colored photos with low saturation. Distinct style with high black contrast and soft camera lens effect.",
    "Fomapan": "Black and white photos with fine/low to moderate grain, highly artistic exposure and angles. Adds very soft lens effect without distortion, dark photo vignette.",
    "Fujicolor": "Color photos with fine/low to moderate grain. Colors are either very oversaturated or slightly desaturated, with entire color hue shifted in a very distinct manner.",
    "Holga": "Color photos with high grain. Colors are either very oversaturated or slightly desaturated. Distinct contrast of black. Often applies photographic vignette.",
    "Instax": "Instant color photos similar to Polaroid but clearer. Near perfect colors, regular saturation, fine/low to medium grain.",
    "Lomography": "Color photos with high grain. Colors are either very oversaturated or slightly desaturated. Distinct contrast of black. Often applies photographic vignette.",
    "Kodachrome": "Color photos with moderate grain. Colors on either colder part of spectrum or regular, with normal or slightly higher saturation.",
    "Rollei": "Mostly black and white photos, sometimes color with fine/low grain. Can be sepia colored or have unusual hues and desaturation. Great for landscapes."
}

composition_styles_info = {
    "Rule of Thirds": "Divides the frame into a 3x3 grid, placing key elements along the lines or at their intersections.",
    "Golden Ratio": "Uses a spiral based on the golden ratio to create a balanced and aesthetically pleasing composition.",
    "Symmetry": "Creates a mirror-like balance in the image, often used for architectural or nature photography.",
    "Leading Lines": "Uses lines within the frame to draw the viewer's eye to the main subject or through the image.",
    "Framing": "Uses elements within the scene to create a frame around the main subject.",
    "Minimalism": "Simplifies the composition to its essential elements, often with a lot of negative space.",
    "Fill the Frame": "The main subject dominates the entire frame, leaving little to no background.",
    "Negative Space": "Uses empty space around the subject to create a sense of simplicity or isolation.",
    "Centered Composition": "Places the main subject in the center of the frame, creating a sense of stability or importance.",
    "Diagonal Lines": "Uses diagonal elements to create a sense of movement or dynamic tension in the image.",
    "Triangular Composition": "Arranges elements in the frame to form a triangle, creating a sense of stability and harmony.",
    "Radial Balance": "Arranges elements in a circular pattern around a central point, creating a sense of movement or completeness."
}

lighting_aspects_info = {
    "Natural light": "Uses available light from the sun or sky, often creating soft, even illumination.",
    "Studio lighting": "Controlled artificial lighting setup, allowing for precise manipulation of light and shadow.",
    "Back light": "Light source behind the subject, creating silhouettes or rim lighting effects.",
    "Split light": "Strong light source at 90-degree angle, lighting one half of the subject while leaving the other in shadow.",
    "Broad light": "Light source at an angle to the subject, producing well-lit photographs with soft to moderate shadows.",
    "Dim light": "Weak or distant light source, creating lower than average brightness and often dramatic images.",
    "Flash photography": "Uses a brief, intense burst of light. Can be fill flash (even lighting) or harsh flash (strong contrasts).",
    "Sunlight": "Direct light from the sun, often creating strong contrasts and warm tones.",
    "Moonlight": "Soft, cool light from the moon, often creating a mysterious or romantic atmosphere.",
    "Spotlight": "Focused beam of light illuminating a specific area, creating high contrast between light and shadow.",
    "High-key lighting": "Bright, even lighting with minimal shadows, creating a light and airy feel.",
    "Low-key lighting": "Predominantly dark tones with selective lighting, creating a moody or dramatic atmosphere.",
    "Rembrandt lighting": "Classic portrait lighting technique creating a triangle of light on the cheek of the subject."
}

special_techniques_info = {
    "Double exposure": "Superimposes two exposures to create a single image, often resulting in a dreamy or surreal effect.",
    "Long exposure": "Uses a long shutter speed to capture motion over time, often creating smooth, blurred effects for moving elements.",
    "Multiple exposure": "Superimposes multiple exposures, multiplying the subject or its key elements across the image.",
    "HDR": "High Dynamic Range imaging, combining multiple exposures to capture a wider range of light and dark tones.",
    "Bokeh effect": "Creates a soft, out-of-focus background, often with circular highlights.",
    "Silhouette": "Captures the outline of a subject against a brighter background, creating a dramatic contrast.",
    "Panning": "Follows a moving subject with the camera, creating a sharp subject with a blurred background.",
    "Light painting": "Uses long exposure and moving light sources to 'paint' with light in the image.",
    "Infrared photography": "Captures light in the infrared spectrum, often resulting in surreal, otherworldly images.",
    "Ultraviolet photography": "Captures light in the ultraviolet spectrum, often revealing hidden patterns or creating a strong violet glow.",
    "Kirlian photography": "High-voltage photographic technique that captures corona discharges around objects, creating a glowing effect.",
    "Thermography": "Captures infrared radiation to create images based on temperature differences, resulting in false-color heat maps.",
    "Astrophotography": "Specialized technique for capturing astronomical objects and celestial events, often resulting in stunning starry backgrounds.",
    "Underwater photography": "Captures images beneath the surface of water, often in pools, seas, or aquariums.",
    "Aerial photography": "Captures images from an elevated position, such as from drones, helicopters, or planes.",
    "Macro photography": "Extreme close-up photography, revealing tiny details not visible to the naked eye."
}

color_effects_info = {
    "Black and white": "Removes all color, leaving only shades of gray.",
    "Sepia": "Reddish-brown monochrome effect, often associated with vintage photography.",
    "Monochrome": "Uses variations of a single color.",
    "Vintage color": "Muted or faded color palette reminiscent of old photographs.",
    "Cross-processed": "Deliberate processing of film in the wrong chemicals, creating unusual color shifts.",
    "Desaturated": "Reduces the intensity of all colors in the image.",
    "Vivid colors": "Increases the saturation and intensity of colors.",
    "Pastel colors": "Soft, pale colors with a light and airy feel.",
    "High contrast": "Emphasizes the difference between light and dark areas in the image.",
    "Low contrast": "Reduces the difference between light and dark areas, creating a softer look.",
    "Color splash": "Converts most of the image to black and white while leaving one or more elements in color."
}

def get_dropdown_choices(info_dict):
    return [f"{key}: {value}" for key, value in info_dict.items()]

def login(username, password):
    if username == USERNAME and password == PASSWORD:
        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(value="Login successful! You can now access the Caption Captain tab.", visible=True)
    else:
        return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(value="Invalid username or password. Please try again.", visible=True)

# Gradio interface
with gr.Blocks(theme="Hev832/Applio", css=css, fill_width=True, fill_height=True) as demo:
    with gr.Tab("Welcome"):
        with gr.Row():
            with gr.Column(scale=2): 
                gr.Markdown(
                    """
                    <img src="https://cdn-uploads.huggingface.co/production/uploads/64740cf7485a7c8e1bd51ac9/LVZnwLV43UUvKu3HORqSs.webp" alt="UDG" width="250" class="centered-image">
                    
                    # 🎨 Underground Digital's Caption Captain: AI-Powered Art Inspiration
                    
                    ## Accelerate Your Creative Workflow with Intelligent Image Analysis
                    
                    This innovative tool empowers Yamamoto's artists to quickly generate descriptive captions,<br>
                    training prompts, and tags from existing artwork, fueling the creative process for GenAI models.
                    
                    ## 🚀 How It Works:
                    
                    1. **Upload Your Inspiration**: Drop in an image (e.g., a charcoal horse picture) that embodies your desired style.
                    2. **Choose Your Output**: Select from descriptive captions, training prompts, and tags.
                    3. **Customize the Results**: Adjust tone, length, and other parameters to fine-tune the output.
                    4. **Generate and Iterate**: Click 'Caption' to analyze your image and use the results to inspire new creations.
                    """
                )

            with gr.Column(scale=1):
                with gr.Row():
                    gr.Markdown(
                        """
                        Login below using the internal<br>
                        username and password to access the full app.<br>
                        
                        Once logged in, a new tab will appear named<br>
                        "Caption Captain" allowing you to access the app.
                        """
                    )

                with gr.Row():
                    username = gr.Textbox(label="Username", placeholder="Enter your username", value="ugd")
                with gr.Row():
                    password = gr.Textbox(label="Password", type="password", placeholder="Enter your password", value="ugd!")
                with gr.Row():
                    login_button = gr.Button("Login", size="sm")
                login_message = gr.Markdown(visible=False)

    caption_captain_tab = gr.Tab("Caption Captain", visible=False)
    with caption_captain_tab:
        with gr.Accordion("How to Use Caption Captain", open=False):
            gr.Markdown("""
            # How to Use Caption Captain
            <img src="https://cdn-uploads.huggingface.co/production/uploads/64740cf7485a7c8e1bd51ac9/Ce_Z478iOXljvpZ_Fr_Y7.png" alt="Captain" width="100" style="max-width: 100%; height: auto;">
                        
            Hello, artist! Let's create amazing captions for your pictures. Here's a comprehensive guide:
            1. **Upload Your Image**: Choose a picture you want to caption and upload it.
            2. **Select Caption Type**: 
               - "Descriptive": Provides a detailed description of the image
               - "Descriptive (Informal)": Gives a casual description of the image
               - "Stable Diffusion Prompt": Creates prompts for AI image generation
               - "MidJourney": Generates prompts specifically for MidJourney AI
               - "Booru tag list": Produces tags in the style of image booru websites
               - "Booru-like tag list": Creates tags similar to booru style, but more flexible
               - "Art Critic": Analyzes the image from an art critic's perspective
               - "Product Listing": Describes the image as if it were a product for sale
               - "Social Media Post": Crafts a caption suitable for social media
               - "Style Prompt": Generates detailed prompts for specific image styles
            3. **Choose Caption Length**:
               - "Any": Lets the AI decide the appropriate length
               - Select from "very short" to "very long"
               - Or specify an exact word count (20 to 260 words)
            4. **Extra Options**: Customize your caption with these choices:
               - Refer to people/characters as a specific name
               - Focus on changeable attributes, not unchangeable characteristics
               - Include details about lighting, camera angle, or image quality
               - Specify technical photography details for photos
               - Control the tone (e.g., keep it PG, avoid ambiguous language)
               - Include or exclude specific elements (e.g., text in the image, aesthetic quality)
               - Determine the level of detail in the description
            5. **Advanced Options for Style Prompt**:
               - Lens Type: Choose from various lens effects (e.g., wide-angle, telephoto)
               - Film Stock: Select different film types for unique color and grain effects
               - Composition Style: Specify how elements are arranged in the image
               - Lighting Aspect: Define the mood through lighting choices
               - Special Technique: Add unique photographic effects
               - Color Effect: Alter the overall color palette of the image
            6. **Generate Your Caption**: Click "Make My Caption!" and see the results!
            ## Tips for Great Captions:
            - Experiment with different caption types to find what works best for your needs
            - Use the extra options to fine-tune the focus and style of your captions
            - For "Style Prompt", explore the advanced options to create highly specific image descriptions
            - If you're not satisfied with a caption, simply generate a new one
            - The "Custom Prompt" option allows for more specific instructions, but use it cautiously
            Remember, Caption Captain is here to inspire and assist your creative process. Have fun exploring all the possibilities!
            """)

        with gr.Row():
            with gr.Column():
                input_image = gr.Image(type="pil", label="Input Image")

                caption_type = gr.Dropdown(
                    choices=["Descriptive", "Descriptive (Informal)", "Stable Diffusion Prompt", "MidJourney", "Booru tag list", "Booru-like tag list", "Art Critic", "Product Listing", "Social Media Post", "Style Prompt"],
                    label="Caption Type",
                    value="Descriptive",
                )

                caption_length = gr.Dropdown(
                    choices=["any", "very short", "short", "medium-length", "long", "very long"] +
                            [str(i) for i in range(20, 261, 10)],
                    label="Caption Length",
                    value="long",
                )

                with gr.Accordion("Extra Options", open=False):
                    extra_options = gr.CheckboxGroup(
                        choices=[
                            "If there is a person/character in the image you must refer to them as {name}.",
                            "Do NOT include information about people/characters that cannot be changed (like ethnicity, gender, etc), but do still include changeable attributes (like hair style).",
                            "Include information about lighting.",
                            "Include information about camera angle.",
                            "Include information about whether there is a watermark or not.",
                            "Include information about whether there are JPEG artifacts or not.",
                            "If it is a photo you MUST include information about what camera was likely used and details such as aperture, shutter speed, ISO, etc.",
                            "Do NOT include anything sexual; keep it PG.",
                            "Do NOT mention the image's resolution.",
                            "You MUST include information about the subjective aesthetic quality of the image from low to very high.",
                            "Include information on the image's composition style, such as leading lines, rule of thirds, or symmetry.",
                            "Do NOT mention any text that is in the image.",
                            "Specify the depth of field and whether the background is in focus or blurred.",
                            "If applicable, mention the likely use of artificial or natural lighting sources.",
                            "Do NOT use any ambiguous language.",
                            "Include whether the image is sfw, suggestive, or nsfw.",
                            "ONLY describe the most important elements of the image."
                        ],
                        label="Select Extra Options"
                    )

                name_input = gr.Textbox(label="Person/Character Name (if applicable)")
                gr.Markdown("**Note:** Name input is only used if an Extra Option is selected that requires it.")

                custom_prompt = gr.Textbox(label="Custom Prompt (optional, will override all other settings)")
                gr.Markdown("**Note:** Alpha Two is not a general instruction follower and will not follow prompts outside its training data well. Use this feature with caution.")

            with gr.Column():
                error_message = gr.Markdown(visible=False)
                output_prompt = gr.Textbox(label="Prompt that was used")
                output_caption = gr.Textbox(label="Generated Caption")

                run_button = gr.Button("Make My Caption!")

                with gr.Accordion("Advanced Options for Style Prompt", open=False, visible=False) as advanced_options:
                    gr.Markdown("### Advanced Options for Style Prompt")
                    lens_type = gr.Dropdown(
                        choices=get_dropdown_choices(lens_types_info),
                        label="Lens Type",
                        info="Select a lens type to define the perspective and field of view of the image."
                    )
                    film_stock = gr.Dropdown(
                        choices=get_dropdown_choices(film_stocks_info),
                        label="Film Stock",
                        info="Choose a film stock to determine the color, grain, and overall look of the image."
                    )
                    composition_style = gr.Dropdown(
                        choices=get_dropdown_choices(composition_styles_info),
                        label="Composition Style",
                        info="Select a composition style to guide the arrangement of elements in the image."
                    )
                    lighting_aspect = gr.Dropdown(
                        choices=get_dropdown_choices(lighting_aspects_info),
                        label="Lighting Aspect",
                        info="Choose a lighting style to define the mood and atmosphere of the image."
                    )
                    special_technique = gr.Dropdown(
                        choices=get_dropdown_choices(special_techniques_info),
                        label="Special Technique",
                        info="Select a special photographic technique to add unique effects to the image."
                    )
                    color_effect = gr.Dropdown(
                        choices=get_dropdown_choices(color_effects_info),
                        label="Color Effect",
                        info="Choose a color effect to alter the overall color palette of the image."
                    )

    def update_style_options(caption_type):
        return gr.update(visible=caption_type == "Style Prompt")

    caption_type.change(update_style_options, inputs=[caption_type], outputs=[advanced_options])
    
    def process_and_handle_errors(input_image, caption_type, caption_length, extra_options, name_input, custom_prompt, lens_type, film_stock, composition_style, lighting_aspect, special_technique, color_effect):
        try:
            prompt, result = stream_chat(input_image, caption_type, caption_length, extra_options, name_input, custom_prompt, lens_type, film_stock, composition_style, lighting_aspect, special_technique, color_effect)
            return gr.update(visible=False), prompt, result
        except Exception as e:
            return gr.update(visible=True, value=f"Error: {str(e)}"), "", ""

    run_button.click(
        fn=process_and_handle_errors,
        inputs=[input_image, caption_type, caption_length, extra_options, name_input, custom_prompt, lens_type, film_stock, composition_style, lighting_aspect, special_technique, color_effect],
        outputs=[error_message, output_prompt, output_caption]
    )

    login_button.click(
        login,
        inputs=[username, password],
        outputs=[caption_captain_tab, username, password, login_message]
    )


    password.submit(
        login,
        inputs=[username, password],
        outputs=[caption_captain_tab, username, password, login_message]
    )

if __name__ == "__main__":
    demo.launch()