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("9em124t2-499968")
TITLE = "<h1><center>JoyCaption Alpha One (2024-09-20a)</center></h1>"
CAPTION_TYPE_MAP = {
	("descriptive", "formal", False, False): ["Write a descriptive caption for this image in a formal tone."],
	("descriptive", "formal", False, True): ["Write a descriptive caption for this image in a formal tone within {word_count} words."],
	("descriptive", "formal", True, False): ["Write a {length} descriptive caption for this image in a formal tone."],
	("descriptive", "informal", False, False): ["Write a descriptive caption for this image in a casual tone."],
	("descriptive", "informal", False, True): ["Write a descriptive caption for this image in a casual tone within {word_count} words."],
	("descriptive", "informal", True, False): ["Write a {length} descriptive caption for this image in a casual tone."],

	("training_prompt", "formal", False, False): ["Write a stable diffusion prompt for this image."],
	("training_prompt", "formal", False, True): ["Write a stable diffusion prompt for this image within {word_count} words."],
	("training_prompt", "formal", True, False): ["Write a {length} stable diffusion prompt for this image."],

	("rng-tags", "formal", False, False): ["Write a list of Booru tags for this image."],
	("rng-tags", "formal", False, True): ["Write a list of Booru tags for this image within {word_count} words."],
	("rng-tags", "formal", True, False): ["Write a {length} list of Booru tags for this image."],

	("style_prompt", "formal", False, False): ["Generate a detailed style prompt for this image, including lens type, film stock, composition notes, lighting aspects, and any special photographic techniques."],
	("style_prompt", "formal", False, True): ["Generate a detailed style prompt for this image within {word_count} words, including lens type, film stock, composition notes, lighting aspects, and any special photographic techniques."],
	("style_prompt", "formal", True, False): ["Generate a {length} detailed style prompt for this image, including lens type, film stock, composition notes, lighting aspects, and any special photographic techniques."],
}

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(MODEL_PATH, use_fast=False)
assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance(tokenizer, PreTrainedTokenizerFast), f"Tokenizer is of type {type(tokenizer)}"

# LLM
print("Loading LLM")
if (CHECKPOINT_PATH / "text_model").exists:
	print("Loading VLM's custom text model")
	text_model = AutoModelForCausalLM.from_pretrained(CHECKPOINT_PATH / "text_model", device_map=0, torch_dtype=torch.bfloat16)
else:
	text_model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map="auto", 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", weights_only=True))
image_adapter.eval()
image_adapter.to("cuda")


def preprocess_image(input_image: Image.Image) -> torch.Tensor:
    """
    Preprocess the input image for the CLIP model.
    """
    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:
    """
    Generate a caption based on the image features and prompt.
    """
    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'))
    embedded_bos = text_model.model.embed_tokens(torch.tensor([[tokenizer.bos_token_id]], device=text_model.device, dtype=torch.int64))
    eot_embed = image_adapter.get_eot_embedding().unsqueeze(0).to(dtype=text_model.dtype)

    inputs_embeds = torch.cat([
        embedded_bos.expand(image_features.shape[0], -1, -1),
        image_features.to(dtype=embedded_bos.dtype),
        prompt_embeds.expand(image_features.shape[0], -1, -1),
        eot_embed.expand(image_features.shape[0], -1, -1),
    ], dim=1)

    input_ids = torch.cat([
        torch.tensor([[tokenizer.bos_token_id]], dtype=torch.long),
        torch.zeros((1, image_features.shape[1]), dtype=torch.long),
        prompt,
        torch.tensor([[tokenizer.convert_tokens_to_ids("<|eot_id|>")]], dtype=torch.long),
    ], dim=1).to('cuda')
    attention_mask = torch.ones_like(input_ids)

    generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_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].strip()

@spaces.GPU()
@torch.no_grad()
def stream_chat(input_image: Image.Image, caption_type: str, caption_tone: str, caption_length: str | int, lens_type: str = "", film_stock: str = "", composition_style: str = "", lighting_aspect: str = "", special_technique: str = "", color_effect: str = "") -> str:
    """
    Generate a caption, training prompt, tags, or a style prompt for image generation based on the input image and parameters.
    """
    # Check if an image has been uploaded
    if input_image is None:
        return "Error: Please upload an image before generating a caption."

    torch.cuda.empty_cache()

    try:
        length = None if caption_length == "any" else caption_length
        if isinstance(length, str):
            length = int(length)
    except ValueError:
        raise ValueError(f"Invalid caption length: {caption_length}")

    if caption_type in ["rng-tags", "training_prompt", "style_prompt"]:
        caption_tone = "formal"

    prompt_key = (caption_type, caption_tone, isinstance(length, str), isinstance(length, int))
    if prompt_key not in CAPTION_TYPE_MAP:
        raise ValueError(f"Invalid caption type: {prompt_key}")

    if caption_type == "style_prompt":
        # For style prompt, we'll create a custom prompt for the LLM
        base_prompt = "Analyze the given image and create a detailed Stable Diffusion prompt for generating a new, creative image inspired by it. "
        base_prompt += "The prompt should describe the main elements, style, and mood of the image, "
        base_prompt += "but also introduce creative variations or enhancements. "
        base_prompt += "Include specific details about the composition, lighting, and overall atmosphere. "
        
        # Add custom settings to the prompt
        if lens_type:
            lens_type_key = lens_type.split(":")[0].strip()
            base_prompt += f"Incorporate the effect of a {lens_type_key} lens ({lens_types_info[lens_type_key]}). "
        if film_stock:
            film_stock_key = film_stock.split(":")[0].strip()
            base_prompt += f"Apply the characteristics of {film_stock_key} film stock ({film_stocks_info[film_stock_key]}). "
        if composition_style:
            composition_style_key = composition_style.split(":")[0].strip()
            base_prompt += f"Use a {composition_style_key} composition style ({composition_styles_info[composition_style_key]}). "
        if lighting_aspect:
            lighting_aspect_key = lighting_aspect.split(":")[0].strip()
            base_prompt += f"Implement {lighting_aspect_key} lighting ({lighting_aspects_info[lighting_aspect_key]}). "
        if special_technique:
            special_technique_key = special_technique.split(":")[0].strip()
            base_prompt += f"Apply the {special_technique_key} technique ({special_techniques_info[special_technique_key]}). "
        if color_effect:
            color_effect_key = color_effect.split(":")[0].strip()
            base_prompt += f"Use a {color_effect_key} color effect ({color_effects_info[color_effect_key]}). "
        
        base_prompt += f"The final prompt should be approximately {length} words long. "
        base_prompt += "Format the output as a single paragraph without numbering or bullet points."

        prompt_str = base_prompt
    else:
        prompt_str = CAPTION_TYPE_MAP[prompt_key][0].format(length=length, word_count=length)

    # Debugging: Print the constructed prompt string
    print(f"Constructed Prompt: {prompt_str}")

    pixel_values = preprocess_image(input_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')

    # Load the model from MODEL_PATH
    text_model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map="auto", torch_dtype=torch.bfloat16)
    text_model.eval()

    # Debugging: Print the prompt string before passing to generate_caption
    print(f"Prompt passed to generate_caption: {prompt_str}")

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

    return caption

css = """
h1, h2, h3, h4, h5, h6, p, li, ul, ol, a {
    text-align: left;
}
.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;
}
"""

# 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/gW2pkMIzSg0REqju2nHsB.png" alt="UDG" alt="UGD Logo" 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, or 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/sDjwwSS4L_atPLP_H5Glv.png" alt="Captain" width="100" style="max-width: 100%; height: auto;">
            
            Hello, artist! Let's make some fun captions for your pictures. Here's how:

            1. **Pick a Picture**: Find a cool picture you want to talk about and upload it.

            2. **Choose What You Want**: 
               - **Caption Type**: 
                 * "Descriptive" tells you what's in the picture
                 * "Training Prompt" helps computers make similar pictures
                 * "RNG-Tags" gives you short words about the picture
                 * "Style Prompt" creates detailed prompts for image generation

            3. **Pick a Style** (for "Descriptive" and "Style Prompt" only):
               - "Formal" sounds like a teacher talking
               - "Informal" sounds like a friend chatting

            4. **Decide How Long**:
               - "Any" lets the computer decide
               - Or pick a size from "very short" to "very long"
               - You can even choose a specific number of words!

            5. **Advanced Options** (for "Style Prompt" only):
               - Choose lens type, film stock, composition, and lighting details

            6. **Make the Caption**: Click the "Make My Caption!" button and watch the magic happen!

            Remember, have fun and be creative with your captions!

            ## Tips for Great Captions:
            - Try different types to see what you like best
            - Experiment with formal and informal tones for fun variations
            - Adjust the length to get just the right amount of detail
            - For "Style Prompt", play with the advanced options for more specific results
            - If you don't like a caption, just click "Make My Caption!" again for a new one
            
            Have a great time captioning your art!
            """)

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

                caption_type = gr.Dropdown(
                    choices=["descriptive", "training_prompt", "rng-tags", "style_prompt"],
                    label="Caption Type",
                    value="descriptive",
                )

                caption_tone = gr.Dropdown(
                    choices=["formal", "informal"],
                    label="Caption Tone",
                    value="formal",
                )

                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="any",
                )

                gr.Markdown("**Note:** Caption tone doesn't affect `rng-tags`, `training_prompt`, and `style_prompt`.")

                
            with gr.Column():
                error_message = gr.Markdown(visible=False)  # Add this line
                output_caption = gr.Textbox(label="Generated Caption")
                run_button = gr.Button("Make My Caption!")

                # Container for advanced options
                with gr.Column(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_tone, caption_length, lens_type, film_stock, composition_style, lighting_aspect, special_technique, color_effect):
        try:
            result = stream_chat(input_image, caption_type, caption_tone, caption_length, lens_type, film_stock, composition_style, lighting_aspect, special_technique, color_effect)
            return gr.update(visible=False), 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_tone, caption_length, lens_type, film_stock, composition_style, lighting_aspect, special_technique, color_effect],
        outputs=[error_message, output_caption]
    )

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

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