Update app.py

app.py

@@ -1,18 +1,15 @@
+import spaces
 import gradio as gr
-from PIL import Image
-import qrcode
+from huggingface_hub import InferenceClient
+from torch import nn
+from transformers import AutoModel, AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM
 from pathlib import Path
-import requests
-import io
-import os
+import torch
+import torch.amp.autocast_mode
 from PIL import Image
-import numpy as np
-import cv2
-from pyzxing import BarCodeReader
-from PIL import ImageOps, ImageEnhance, ImageFilter
-from huggingface_hub import hf_hub_download, snapshot_download
-from PIL import ImageEnhance
-import replicate
+import os
+import torchvision.transforms.functional as TVF
+
 from dotenv import load_dotenv
 
 # Load environment variables from .env file
@@ -20,348 +17,438 @@ load_dotenv()
 
 USERNAME = os.getenv("USERNAME")
 PASSWORD = os.getenv("PASSWORD")
-REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")
-
-# Set the Replicate API token
-os.environ["REPLICATE_API_TOKEN"] = REPLICATE_API_TOKEN
-
-qrcode_generator = qrcode.QRCode(
-    version=1,
-    error_correction=qrcode.ERROR_CORRECT_H,
-    box_size=10,
-    border=4,
-)
-
 
-# Define available models
-CONTROLNET_MODELS = {
-    "QR Code Monster": "monster-labs/control_v1p_sd15_qrcode_monster/v2/",
-    "QR Code": "DionTimmer/controlnet_qrcode-control_v1p_sd15",
-    # Add more ControlNet models here
+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."],
 }
 
-DIFFUSION_MODELS = {
-    "GhostMix": "digiplay/GhostMixV1.2VAE",
-    "Stable v1.5": "Jiali/stable-diffusion-1.5",
-    # Add more diffusion models here
-}
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
 
-# Global variables to store loaded models
-loaded_controlnet = None
-loaded_pipe = None
 
-#   def load_models_on_launch():
-#    global loaded_controlnet, loaded_pipe
-#    print("Loading models on launch...")
-    
-# Download the main repository
-#    main_repo_path = snapshot_download("monster-labs/control_v1p_sd15_qrcode_monster")
-    
-# Construct the path to the subfolder
-#    controlnet_path = os.path.join(main_repo_path, "v2")
-
-#    loaded_controlnet = ControlNetModel.from_pretrained(
-#        controlnet_path,
-#        torch_dtype=torch.float16
-#    ).to("mps")
-
-#    diffusion_path = snapshot_download(DIFFUSION_MODELS["GhostMix"])
-#    loaded_pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-#        diffusion_path,
-#        controlnet=loaded_controlnet,
-#        torch_dtype=torch.float16,
-#        safety_checker=None,
-#    ).to("mps")
-#    print("Models loaded successfully!")
-
-# Modify the load_models function to use global variables
-#def load_models(controlnet_model, diffusion_model):
-#    global loaded_controlnet, loaded_pipe
-#    if loaded_controlnet is None or loaded_pipe is None:
-#        load_models_on_launch()
-#    return loaded_pipe
-
-# Add new functions for image adjustments
-def adjust_image(image, brightness, contrast, saturation):
-    if image is None:
-        return None
-    
-    img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
-    
-    if brightness != 1:
-        img = ImageEnhance.Brightness(img).enhance(brightness)
-    if contrast != 1:
-        img = ImageEnhance.Contrast(img).enhance(contrast)
-    if saturation != 1:
-        img = ImageEnhance.Color(img).enhance(saturation)
-    
-    return np.array(img)
-
-def resize_for_condition_image(input_image: Image.Image, resolution: int):
-    input_image = input_image.convert("RGB")
-    W, H = input_image.size
-    k = float(resolution) / min(H, W)
-    H *= k
-    W *= k
-    H = int(round(H / 64.0)) * 64
-    W = int(round(W / 64.0)) * 64
-    img = input_image.resize((W, H), resample=Image.LANCZOS)
-    return img
-
-
-#   SAMPLER_MAP = {
-#    "DPM++ Karras SDE": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True, algorithm_type="sde-dpmsolver++"),
-#    "DPM++ Karras": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True),
-#    "Heun": lambda config: HeunDiscreteScheduler.from_config(config),
-#    "Euler": lambda config: EulerDiscreteScheduler.from_config(config),
-#    "DDIM": lambda config: DDIMScheduler.from_config(config),
-#    "DEIS": lambda config: DEISMultistepScheduler.from_config(config),
-#}
-
-def scan_qr_code(image):
-    # Convert gradio image to PIL Image if necessary
-    if isinstance(image, np.ndarray):
-        image = Image.fromarray(image)
-    
-    # Convert to grayscale
-    gray_image = image.convert('L')
-    
-    # Convert to numpy array
-    np_image = np.array(gray_image)
-    
-    # Method 1: Using qrcode library
-    try:
-        qr = qrcode.QRCode()
-        qr.add_data('')
-        qr.decode(gray_image)
-        return qr.data.decode('utf-8')
-    except Exception:
-        pass
-    
-    # Method 2: Using OpenCV
-    try:
-        qr_detector = cv2.QRCodeDetector()
-        retval, decoded_info, points, straight_qrcode = qr_detector.detectAndDecodeMulti(np_image)
-        if retval:
-            return decoded_info[0]
-    except Exception:
-        pass
-    
-    # Method 3: Fallback to zxing-cpp
+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:
-        reader = BarCodeReader()
-        results = reader.decode(np_image)
-        if results:
-            return results[0].parsed
-    except Exception:
-        pass
-    
-    return None
-
-def invert_image(image):
-    if image is None:
-        return None
-    if isinstance(image, np.ndarray):
-        return 255 - image
-    elif isinstance(image, Image.Image):
-        return ImageOps.invert(image.convert('RGB'))
+        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:
-        raise ValueError("Unsupported image type")
-
-def invert_displayed_image(image):
-    if image is None:
-        return None
-    inverted = invert_image(image)
-    if isinstance(inverted, np.ndarray):
-        return Image.fromarray(inverted)
-    return inverted
-
-
-#@spaces.GPU()
-def inference(
-    qr_code_content: str,
-    prompt: str,
-    negative_prompt: str,
-    guidance_scale: float = 9.0,
-    qr_conditioning_scale: float = 1.47,
-    num_inference_steps: int = 20,
-    seed: int = -1,
-    image_resolution: int = 512,
-    scheduler: str = "K_EULER",
-    eta: float = 0.0,
-    num_outputs: int = 1,
-    low_threshold: int = 100,
-    high_threshold: int = 200,
-    guess_mode: bool = False,
-    disable_safety_check: bool = False,
-):
-    try:
-        progress = gr.Progress()
-        progress(0, desc="Generating QR code...")
-
-        # Generate QR code image
-        qr = qrcode.QRCode(
-            version=1,
-            error_correction=qrcode.constants.ERROR_CORRECT_H,
-            box_size=10,
-            border=4,
-        )
-        qr.add_data(qr_code_content)
-        qr.make(fit=True)
-        qr_image = qr.make_image(fill_color="black", back_color="white")
-
-        # Save QR code image to a temporary file
-        temp_qr_path = "temp_qr.png"
-        qr_image.save(temp_qr_path)
-
-        progress(0.3, desc="Running inference...")
-
-        # Ensure num_outputs is within the allowed range
-        num_outputs = max(1, min(num_outputs, 10))
-
-        # Ensure seed is an integer and not null
-        seed = int(seed) if seed != -1 else None
-
-        # Ensure high_threshold is at least 1
-        high_threshold = max(1, int(high_threshold))
-
-        # Prepare the input dictionary
-        input_dict = {
-            "prompt": prompt,
-            "qr_image": open(temp_qr_path, "rb"),
-            "negative_prompt": negative_prompt,
-            "guidance_scale": float(guidance_scale),
-            "qr_conditioning_scale": float(qr_conditioning_scale),
-            "num_inference_steps": int(num_inference_steps),
-            "image_resolution": int(image_resolution),
-            "scheduler": scheduler,
-            "eta": float(eta),
-            "num_outputs": num_outputs,
-            "low_threshold": int(low_threshold),
-            "high_threshold": high_threshold,
-            "guess_mode": guess_mode,
-            "disable_safety_check": disable_safety_check,
-        }
-
-        # Only add seed to input_dict if it's not None
-        if seed is not None:
-            input_dict["seed"] = seed
-
-        # Run inference using Replicate API
-        output = replicate.run(
-            "anotherjesse/multi-control:76d8414a702e66c84fe2e6e9c8cbdc12e53f950f255aae9ffa5caa7873b12de0",
-            input=input_dict
-        )
-
-        progress(0.9, desc="Processing results...")
-
-        # Download the generated image
-        response = requests.get(output[0])
-        img = Image.open(io.BytesIO(response.content))
-
-        # Clean up temporary file
-        os.remove(temp_qr_path)
-
-        progress(1.0, desc="Done!")
-        return img, seed if seed is not None else -1
-    except Exception as e:
-        print(f"Error in inference: {str(e)}")
-        return Image.new('RGB', (512, 512), color='white'), -1
-
-
-
-def invert_init_image_display(image):
-    if image is None:
-        return None
-    inverted = invert_image(image)
-    if isinstance(inverted, np.ndarray):
-        return Image.fromarray(inverted)
-    return inverted
-
-def adjust_color_balance(image, r, g, b):
-    # Convert image to RGB if it's not already
-    image = image.convert('RGB')
-    
-    # Split the image into its RGB channels
-    r_channel, g_channel, b_channel = image.split()
-    
-    # Adjust each channel
-    r_channel = r_channel.point(lambda i: i + (i * r))
-    g_channel = g_channel.point(lambda i: i + (i * g))
-    b_channel = b_channel.point(lambda i: i + (i * b))
-    
-    # Merge the channels back
-    return Image.merge('RGB', (r_channel, g_channel, b_channel))
+        prompt_str = CAPTION_TYPE_MAP[prompt_key][0].format(length=length, word_count=length)
 
-def apply_qr_overlay(image, original_qr, overlay, opacity):
-    if not overlay or original_qr is None:
-        return image
-    
-    # Resize original QR to match the generated image
-    original_qr = original_qr.resize(image.size)
-    
-    # Create a new image blending the generated image and the QR code
-    return Image.blend(image, original_qr, opacity)
+    # Debugging: Print the constructed prompt string
+    print(f"Constructed Prompt: {prompt_str}")
 
-def apply_edge_enhancement(image, strength):
-    if strength == 0:
-        return image
-    
-    # Apply edge enhancement
-    enhanced = image.filter(ImageFilter.EDGE_ENHANCE)
-    
-    # Blend the original and enhanced images based on strength
-    return Image.blend(image, enhanced, strength / 5.0)
+    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, .centered-image {
-    text-align: center;
+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 {
-    margin-left: auto;
-    margin-right: auto;
-    display: table;
+    padding-left: 20px;
 }
-.centered-image {
-    max-width: 100%;
-    height: auto;
+.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 QR Code Art Generator 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)
-    
-# Add login elements to the Gradio interface
-with gr.Blocks(theme='Hev832/Applio', css=css, fill_width=True, fill_height=True) as blocks:
-    generated_images = gr.State([])
 
+# 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):  
+            with gr.Column(scale=2): 
                 gr.Markdown(
                     """
-                    <img src="https://cdn-uploads.huggingface.co/production/uploads/64740cf7485a7c8e1bd51ac9/8lHjpId7-JDalHq1JByPE.webp" alt="Yamamoto Logo" class="centered-image">
-                    
-                    # 🎨 Yamamoto QR Code Art Generator
-                
-                    ## Transform Your QR Codes into Brand Masterpieces
-                    This cutting-edge tool empowers our creative team to craft visually stunning,<br> 
-                    on-brand QR codes that perfectly blend functionality with artistic expression.
+                    <img src="https://cdn-uploads.huggingface.co/production/uploads/64740cf7485a7c8e1bd51ac9/LVZnwLV43UUvKu3HORqSs.webp" alt="UDG" width="250" style="max-width: 100%; height: auto; 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. **Enter Your QR Code Content**: Start by inputting the URL or text for your QR code.
-                    2. **Craft Your Prompt**: Describe the artistic style or theme you envision for your QR code.
-                    3. **Fine-tune with Advanced Settings**: Adjust parameters to perfect your creation (see tips below).
-                    4. **Generate and Iterate**: Click 'Run' to create your art, then refine as needed.
+                    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(
@@ -382,479 +469,133 @@ with gr.Blocks(theme='Hev832/Applio', css=css, fill_width=True, fill_height=True
                     login_button = gr.Button("Login", size="sm")
                 login_message = gr.Markdown(visible=False)
 
+    with gr.Tab("Caption Captain") as app_container:
+        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 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.Tab("QR Code Art Generator", visible=False) as app_container:
         with gr.Row():
             with gr.Column():
-                qr_code_content = gr.Textbox(
-                    label="QR Code Content",
-                    placeholder="Enter URL or text for your QR code",
-                    info="This is what your QR code will link to or display when scanned.",
-                    value="https://theunderground.digital/",
-                    lines=1,
-                )
+                input_image = gr.Image(type="pil", label="Input Image")
 
-                prompt = gr.Textbox(
-                    label="Artistic Prompt",
-                    placeholder="Describe the style or theme for your QR code art (For best results, keep the prompt to 75 characters or less as seen below)",
-                    value="A high-res, photo-realistic minimalist rendering of Mount Fuji as a sharp, semi-realistic silhouette on the horizon. The mountain conveys strength and motion with clean, crisp lines and natural flow. Features detailed snow textures, subtle ridge highlights, and a powerful yet serene atmosphere. Emphasizes strength with clarity and precision in texture and light.",
-                    info="Describe the style or theme for your QR code art (For best results, keep the prompt to 75 characters or less as seen in the example)",
-                    lines=8,
+                caption_type = gr.Dropdown(
+                    choices=["descriptive", "training_prompt", "rng-tags", "style_prompt"],
+                    label="Caption Type",
+                    value="descriptive",
                 )
-                negative_prompt = gr.Textbox(
-                    label="Elements to Avoid",
-                    placeholder="Describe what you don't want in the image",
-                    value="ugly, disfigured, low quality, blurry, nsfw, bad_pictures, poorly drawn, distorted, overexposed, flat shading, bad proportions, deformed, pixelated, messy details, lack of contrast, unrealistic textures, bad anatomy, rough edges, low resolution",
-                    info="List elements or styles you want to avoid in your QR code art.",
-                    lines=4,
+
+                caption_tone = gr.Dropdown(
+                    choices=["formal", "informal"],
+                    label="Caption Tone",
+                    value="formal",
                 )
 
-                run_btn = gr.Button("🎨 Create Your QR Art", variant="primary")  
+                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",
+                )
 
-                with gr.Accordion(label="Needs Some Prompting Help?", open=False, visible=True):
-                    gr.Markdown(
-                        """
-                        ## 🌟 Tips for Spectacular Results:
-                        - Use concise details in your prompt to help the AI understand your vision.
-                        - Use negative prompts to avoid unwanted elements in your image.
-                        - Experiment with different ControlNet models and diffusion models to find the best combination for your prompt.
-
-                        ## 🎭 Prompt Ideas to Spark Your Creativity:
-                        - "A serene Japanese garden with cherry blossoms and a koi pond"
-                        - "A futuristic cityscape with neon lights and flying cars"
-                        - "An abstract painting with swirling colors and geometric shapes"
-                        - "A vintage-style travel poster featuring iconic landmarks"
-
-                        Remember, the magic lies in the details of your prompt and the fine-tuning of your settings. 
-                        Happy creating!
-                        """
-                    )
+                gr.Markdown("**Note:** Caption tone doesn't affect `rng-tags`, `training_prompt`, and `style_prompt`.")
 
-                with gr.Accordion("Set Custom QR Code Colors", open=False, visible=False):
-                    bg_color = gr.ColorPicker(
-                        label="Background Color", 
-                        value="#FFFFFF",
-                        info="Choose the background color for the QR code"
-                    )
-                    qr_color = gr.ColorPicker(
-                        label="QR Code Color", 
-                        value="#000000",
-                        info="Choose the color for the QR code pattern"
-                    )
-                    invert_final_image = gr.Checkbox(
-                        label="Invert Final Image", 
-                        value=False,
-                        info="Check this to invert the colors of the final image",
-                        visible=False,
-                    )
-                with gr.Accordion("AI Model Selection", open=False, visible=False):
-                    controlnet_model_dropdown = gr.Dropdown(
-                        choices=list(CONTROLNET_MODELS.keys()),
-                        value="QR Code Monster",
-                        label="ControlNet Model",
-                        info="Select the ControlNet model for QR code generation"
-                    )
-                    diffusion_model_dropdown = gr.Dropdown(
-                        choices=list(DIFFUSION_MODELS.keys()),
-                        value="GhostMix",
-                        label="Diffusion Model",
-                        info="Select the main diffusion model for image generation"
-                    )
                 
-            
-                with gr.Accordion(label="QR Code Image (Optional)", open=False, visible=False):
-                    qr_code_image = gr.Image(
-                        label="QR Code Image (Optional). Leave blank to automatically generate QR code",
-                        type="pil",
-                    )
-            
             with gr.Column():
-                gr.Markdown("### Your Generated QR Code Art")
-                result_image = gr.Image(
-                    label="Your Artistic QR Code", 
-                    show_download_button=True, 
-                    show_fullscreen_button=True, 
-                    container=False
-                )
-                gr.Markdown("💾 Right-click and save the image to download your QR code art.")
-
-                scan_button = gr.Button("Verify QR Code Works", visible=False)
-                scan_result = gr.Textbox(label="Validation Result of QR Code", interactive=False, visible=False)
-                used_seed = gr.Number(label="Seed Used", interactive=False)
-
-                with gr.Accordion(label="Use Your Own Image as a Reference", open=True, visible=True) as init_image_acc:
-                    init_image = gr.Image(label="Reference Image", type="pil")
-                    with gr.Row():
-                        use_qr_code_as_init_image = gr.Checkbox(
-                            label="Uncheck to use your own image for generation", 
-                            value=True, 
-                            interactive=True,
-                            info="Allows you to use your own image for generation, otherwise a generic QR Code is created automatically as the base image"
-                        )
-                        reference_image_strength = gr.Slider(
-                            minimum=0.0,
-                            maximum=5.0,
-                            step=0.05,
-                            value=0.6,
-                            label="Reference Image Influence",
-                            info="Controls how much the reference image influences the final result (0 = ignore, 5 = copy exactly)",
-                            visible=False
-                        )
-                        invert_init_image_button = gr.Button("Invert Init Image", size="sm", visible=False)
-
-    with gr.Tab("Advanced Settings"):
-        with gr.Accordion("Advanced Art Controls", open=True):
-            with gr.Row():
-                qr_conditioning_scale = gr.Slider(
-                    minimum=0.0,
-                    maximum=5.0,
-                    step=0.01,
-                    value=1.47,
-                    label="QR Code Visibility",
-                )
-                with gr.Accordion("QR Code Visibility Explanation", open=False):
-                    gr.Markdown(
-                        """
-                        **QR Code Visibility** controls how prominent the QR code is in the final image:
-                        
-                        - **Low (0.0-1.0)**: QR code blends more with the art, potentially harder to scan.
-                        - **Medium (1.0-3.0)**: Balanced visibility, usually scannable while maintaining artistic quality.
-                        - **High (3.0-5.0)**: QR code stands out more, easier to scan but less artistic.
-                        
-                        Start with 1.47 for a good balance between art and functionality.
-                        """
+                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."
                     )
-            
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    minimum=0.1,
-                    maximum=30.0,
-                    step=0.1,
-                    value=9.0,
-                    label="Prompt Adherence",
-                )
-                with gr.Accordion("Prompt Adherence Explanation", open=False):
-                    gr.Markdown(
-                        """
-                        **Prompt Adherence** determines how closely the AI follows your prompt:
-                        
-                        - **Low (0.1-5.0)**: More creative freedom, may deviate from prompt.
-                        - **Medium (5.0-15.0)**: Balanced between prompt and AI creativity.
-                        - **High (15.0-30.0)**: Strictly follows the prompt, less creative freedom.
-                        
-                        A value of 9.0 provides a good balance between creativity and prompt adherence.
-                        """
-                    )
-            
-            with gr.Row():
-                num_inference_steps = gr.Slider(
-                    minimum=1,
-                    maximum=100,
-                    step=1,
-                    value=20,
-                    label="Generation Steps",
-                )
-                with gr.Accordion("Generation Steps Explanation", open=False):
-                    gr.Markdown(
-                        """
-                        **Generation Steps** affects the detail and quality of the generated image:
-                        
-                        - **Low (1-10)**: Faster generation, less detailed results.
-                        - **Medium (11-30)**: Good balance between speed and quality.
-                        - **High (31-100)**: More detailed results, slower generation.
-                        
-                        20 steps is a good starting point for most generations.
-                        """
+                    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."
                     )
-            
-            with gr.Row():
-                image_resolution = gr.Slider(
-                    minimum=256,
-                    maximum=1024,
-                    step=64,
-                    value=512,
-                    label="Image Resolution",
-                )
-                with gr.Accordion("Image Resolution Explanation", open=False):
-                    gr.Markdown(
-                        """
-                        **Image Resolution** determines the size and detail of the generated image:
-                        
-                        - **Low (256-384)**: Faster generation, less detailed.
-                        - **Medium (512-768)**: Good balance of detail and generation time.
-                        - **High (832-1024)**: More detailed, slower generation.
-                        
-                        512x512 is a good default for most use cases.
-                        """
-                    )
-            
-            with gr.Row():
-                seed = gr.Slider(
-                    minimum=-1,
-                    maximum=9999999999,
-                    step=1,
-                    value=-1,
-                    label="Generation Seed",
-                )
-                with gr.Accordion("Generation Seed Explanation", open=False):
-                    gr.Markdown(
-                        """
-                        **Generation Seed** controls the randomness of the generation:
-                        
-                        - **-1**: Random seed each time, producing different results.
-                        - **Any positive number**: Consistent results for the same inputs.
-                        
-                        Use -1 to explore various designs, or set a specific seed to recreate a particular result.
-                        """
-                    )
-            
-            with gr.Row():
-                scheduler = gr.Dropdown(
-                    choices=["DDIM", "K_EULER", "DPMSolverMultistep", "K_EULER_ANCESTRAL", "PNDM", "KLMS"],
-                    value="K_EULER",
-                    label="Sampling Method",
-                )
-                with gr.Accordion("Sampling Method Explanation", open=False):
-                    gr.Markdown(
-                        """
-                        **Sampling Method** affects the image generation process:
-                        
-                        - **K_EULER**: Good balance of speed and quality.
-                        - **DDIM**: Can produce sharper results but may be slower.
-                        - **DPMSolverMultistep**: Often produces high-quality results.
-                        - **K_EULER_ANCESTRAL**: Can introduce more variations.
-                        - **PNDM**: Another quality-focused option.
-                        - **KLMS**: Can produce smooth results.
-                        
-                        Experiment with different methods to find what works best for your specific prompts.
-                        """
-                    )
-
-            with gr.Row():
-                eta = gr.Slider(
-                    minimum=0.0,
-                    maximum=1.0,
-                    step=0.01,
-                    value=0.0,
-                    label="ETA (Noise Level)",
-                )
-                with gr.Accordion("ETA Explanation", open=False):
-                    gr.Markdown(
-                        """
-                        **ETA (Noise Level)** controls the amount of noise in the generation process:
-                        
-                        - **0.0**: No added noise, more deterministic results.
-                        - **0.1-0.5**: Slight variations in output.
-                        - **0.6-1.0**: More variations, potentially more creative results.
-                        
-                        Start with 0.0 and increase if you want more variation in your outputs.
-                        """
+                    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."
                     )
-
-            with gr.Row():
-                low_threshold = gr.Slider(
-                    minimum=1,
-                    maximum=255,
-                    step=1,
-                    value=100,
-                    label="Edge Detection Low Threshold",
-                )
-                high_threshold = gr.Slider(
-                    minimum=1,
-                    maximum=255,
-                    step=1,
-                    value=200,
-                    label="Edge Detection High Threshold",
-                )
-                with gr.Accordion("Edge Detection Thresholds Explanation", open=False):
-                    gr.Markdown(
-                        """
-                        **Edge Detection Thresholds** affect how the QR code edges are processed:
-                        
-                        - **Low Threshold**: Lower values detect more edges, higher values fewer.
-                        - **High Threshold**: Determines which edges are strong. Higher values result in fewer strong edges.
-                        
-                        Default values (100, 200) work well for most QR codes. Adjust if you need more or less edge definition.
-                        """
+                    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."
                     )
-
-            with gr.Row():
-                guess_mode = gr.Checkbox(
-                    label="Guess Mode",
-                    value=False,
-                )
-                with gr.Accordion("Guess Mode Explanation", open=False):
-                    gr.Markdown(
-                        """
-                        **Guess Mode**, when enabled, allows the AI to interpret the input image more freely:
-                        
-                        - **Unchecked**: AI follows the QR code structure more strictly.
-                        - **Checked**: AI has more freedom to interpret the input, potentially leading to more creative results.
-                        
-                        Use this if you want more artistic interpretations of your QR code.
-                        """
+                    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."
                     )
-
-            with gr.Row():
-                disable_safety_check = gr.Checkbox(
-                    label="Disable Safety Check",
-                    value=False,
-                )
-                with gr.Accordion("Safety Check Explanation", open=False):
-                    gr.Markdown(
-                        """
-                        **Disable Safety Check** removes content filtering from the generation process:
-                        
-                        - **Unchecked**: Normal content filtering applied.
-                        - **Checked**: No content filtering, may produce unexpected or inappropriate results.
-                        
-                        Use with caution and only if necessary for your specific use case.
-                        """
+                    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."
                     )
-    with gr.Tab("Image Editing"):
-        with gr.Column():
-            image_selector = gr.Dropdown(label="Select Image to Edit", choices=[], interactive=True, visible=False)
-            image_to_edit = gr.Image(label="Your Artistic QR Code", show_download_button=True, show_fullscreen_button=True, container=True)
-
-            with gr.Row():
-                qr_overlay = gr.Checkbox(label="Overlay Original QR Code", value=False, visible=False)
-                qr_opacity = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="QR Overlay Opacity", visible=False)
-                edge_enhance = gr.Slider(minimum=0.0, maximum=5.0, step=0.1, value=0.0, label="Edge Enhancement", visible=False)
-
-            with gr.Row():
-                red_balance = gr.Slider(minimum=-1.0, maximum=1.0, step=0.1, value=0.0, label="Red Balance")
-                green_balance = gr.Slider(minimum=-1.0, maximum=1.0, step=0.1, value=0.0, label="Green Balance")
-                blue_balance = gr.Slider(minimum=-1.0, maximum=1.0, step=0.1, value=0.0, label="Blue Balance")
-
-
-            with gr.Row():
-                brightness = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Brightness")
-                contrast = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Contrast")
-                saturation = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Saturation")
-            with gr.Row():
-                invert_button = gr.Button("Invert Image", size="sm")
-
-            with gr.Row():
-                edited_image = gr.Image(label="Edited QR Code", show_download_button=True, show_fullscreen_button=True, visible=False)
-                scan_button = gr.Button("Verify QR Code Works", size="sm", visible=False)
-                scan_result = gr.Textbox(label="Validation Result of QR Code", interactive=False, visible=False)
-                
-            used_seed = gr.Number(label="Seed Used", interactive=False)
-            
-            gr.Markdown(
-                """
-                ### 🔍 Analyzing Your Creation
-                - Is the QR code scannable? Check with your phone camera to see if it can scan it.
-                - If not scannable, use the Brightness, Contrast, and Saturation sliders to optimize the QR code for scanning.
-                - Does the art style match your prompt? If not, try adjusting the 'Prompt Adherence'.
-                - Want more artistic flair? Increase the 'Artistic Freedom'.
-                - Need a clearer QR code? Raise the 'QR Code Visibility'.
-                """
-            )
-
-    def scan_and_display(image):
-        if image is None:
-            return "No image to scan"
-        
-        scanned_text = scan_qr_code(image)
-        if scanned_text:
-            return f"Scanned successfully: {scanned_text}"
-        else:
-            return "Failed to scan QR code. Try adjusting the settings for better visibility."
-
-    def invert_displayed_image(image):
-        if image is None:
-            return None
-        return invert_image(image)
-
-    scan_button.click(
-        scan_and_display,
-        inputs=[result_image],
-        outputs=[scan_result]
-    )
-
-    invert_button.click(
-        invert_displayed_image,
-        inputs=[result_image],
-        outputs=[result_image]
-    )
-
-    invert_init_image_button.click(
-        invert_init_image_display,
-        inputs=[init_image],
-        outputs=[init_image]
-    )
-
-    brightness.change(
-        adjust_image,
-        inputs=[result_image, brightness, contrast, saturation],
-        outputs=[result_image]
-    )
-    contrast.change(
-        adjust_image,
-        inputs=[result_image, brightness, contrast, saturation],
-        outputs=[result_image]
-    )
-    saturation.change(
-        adjust_image,
-        inputs=[result_image, brightness, contrast, saturation],
-        outputs=[result_image]
-    )
-
-    # Add logic to show/hide the reference_image_strength slider
-    def update_reference_image_strength_visibility(init_image, use_qr_code_as_init_image):
-        return gr.update(visible=init_image is not None and not use_qr_code_as_init_image)
-
-    init_image.change(
-        update_reference_image_strength_visibility,
-        inputs=[init_image, use_qr_code_as_init_image],
-        outputs=[reference_image_strength]
-    )
 
-    use_qr_code_as_init_image.change(
-        update_reference_image_strength_visibility,
-        inputs=[init_image, use_qr_code_as_init_image],
-        outputs=[reference_image_strength]
-    )
-
-    run_btn.click(
-        fn=inference,
-        inputs=[
-            qr_code_content,
-            prompt,
-            negative_prompt,
-            guidance_scale,
-            qr_conditioning_scale,
-            num_inference_steps,
-            seed,
-            image_resolution,
-            scheduler,
-            eta,
-            low_threshold,
-            high_threshold,
-            guess_mode,
-            disable_safety_check,
-        ],
-        outputs=[result_image, used_seed],
-        concurrency_limit=20
-    )
+    def update_style_options(caption_type):
+        return gr.update(visible=caption_type == "style_prompt")
 
-    # Define login button click behavior
-    login_button.click(
-        login,
-        inputs=[username, password],
-        outputs=[app_container, login_message, login_button, login_message]
-    )
-
-    # Define password textbox submit behavior
-    password.submit(
-        login,
-        inputs=[username, password],
-        outputs=[app_container, login_message, login_button, login_message]
+    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]
     )
 
-# Load models on launch
-#load_models_on_launch()
 
-blocks.queue(max_size=20)
-blocks.launch(share=False, show_api=False)
+if __name__ == "__main__":
+    demo.launch()