import torch from transformers import AutoModelForImageTextToText, AutoProcessor from qwen_vl_utils import process_vision_info from PIL import Image from typing import List, Dict import json from opencc import OpenCC import warnings class CaptionGenerationManager: """Caption generation using Vision-Language Models (supports Qwen2.5-VL, Qwen3-VL, etc.)""" def __init__(self, model_name: str = "Qwen/Qwen2.5-VL-7B-Instruct"): """ Args: model_name: Vision-Language model name, e.g.: - "Qwen/Qwen2.5-VL-7B-Instruct" (default) - "Qwen/Qwen3-VL-8B-Instruct" (2025 latest) """ print(f"Loading Vision-Language Model: {model_name}...") # Suppress processor warning warnings.filterwarnings("ignore", category=FutureWarning, module="transformers") # Use Auto* classes for flexibility (supports Qwen2.5-VL, Qwen3-VL, etc.) self.processor = AutoProcessor.from_pretrained(model_name, use_fast=False) self.model = AutoModelForImageTextToText.from_pretrained( model_name, dtype=torch.bfloat16, # Changed from torch_dtype to dtype device_map="auto" ) # Simplified Chinese to Traditional Chinese converter self.cc = OpenCC('s2t') # Simplified to Traditional self.generation_config = { 'temperature': 0.7, 'top_p': 0.9, 'max_new_tokens': 300, # Increased from 200 to prevent truncation 'repetition_penalty': 1.1 } # Platform-specific templates self.platform_templates = { 'instagram': { 'style': 'storytelling, aesthetic', 'emoji_count': '2-3', 'hashtag_count': '8-10', 'min_length': 120, # Increased for richer content 'max_length': 220, # Allow more detailed descriptions 'features': ['call-to-action', 'question', 'relatable'] }, 'tiktok': { 'style': 'brief, punchy', 'emoji_count': '1-2', 'hashtag_count': '5-8', 'min_length': 60, 'max_length': 120, 'features': ['trending', 'POV', 'relatable'] }, 'xiaohongshu': { 'style': 'structured, informative, detailed', 'emoji_count': '5-8', 'hashtag_count': '8-12', 'min_length': 180, 'max_length': 500, 'features': ['tips', 'bullets', 'sharing-tone'] } } print(f"✓ {model_name.split('/')[-1]} loaded successfully (using Auto* classes for flexibility)") def construct_prompt(self, analysis_results: Dict, platform: str = 'instagram', language: str = 'zh') -> str: """Construct prompt with language support ensuring consistency Args: language: 'zh' (Traditional Chinese), 'en' (English), 'zh-en' (Bilingual) """ platform_config = self.platform_templates.get(platform, self.platform_templates['instagram']) # Language-specific instructions language_instructions = { 'zh': '請使用繁體中文生成標題和標籤。語言要自然流暢，符合華語社群媒體的表達習慣。避免使用簡體字。當偵測到品牌時，必須在標題中提及品牌名稱。', 'en': '''🚨 CRITICAL LANGUAGE REQUIREMENT 🚨 Generate captions and hashtags EXCLUSIVELY in English. - NEVER use Chinese characters (Traditional or Simplified) - NEVER mix languages - Use natural, engaging language suitable for international social media - When brands are detected, mention them naturally in English - All text output must be 100% English only This is MANDATORY and NON-NEGOTIABLE.''', 'zh-en': '''生成雙語內容：標題使用繁體中文，同時提供英文翻譯。標籤混合使用中英文以擴大觸及範圍。當偵測到品牌時，必須在標題中提及品牌名稱。 🚨 重要：雙語一致性要求 🚨 - 中文和英文必須表達相同的核心意義 - 允許表達方式的差異（形容詞、語法不同） - 但整體訊息、語氣、品牌提及必須一致 - 兩種語言都要朝同一方向詮釋內容''' } system_instruction = f"""You are a professional social media content strategist. {language_instructions.get(language, language_instructions['zh'])} Target platform: {platform} Content style: Authentic, creative, and optimized for engagement. CRITICAL RULE: Never include hashtags (symbols starting with #) in the caption text. Hashtags must only appear in the separate 'hashtags' array.""" # Extract analysis context objects = analysis_results.get('detections', []) brands = analysis_results.get('brands', []) scene_info = analysis_results.get('scene_analysis', {}) composition = analysis_results.get('composition', {}) # FIXED: Get fused lighting from scene_info (it's been updated by DetectionFusionManager) lighting = scene_info.get('lighting', {}).get('top', 'natural light') lighting_confidence = scene_info.get('lighting', {}).get('confidence', 0.7) # Provide explicit Chinese translations to ensure consistency lighting_translations_zh = { 'soft diffused light': '柔和漫射光', 'overcast atmosphere': '陰天氛圍', 'natural daylight': '自然日光', 'warm ambient light': '溫暖環境光', 'evening light': '傍晚光線', 'bright sunlight': '明亮陽光', 'golden hour': '金黃時刻', 'blue hour': '藍調時刻' } # Get appropriate lighting description based on language if language == 'zh': lighting_zh = lighting_translations_zh.get(lighting, lighting) lighting_display = lighting_zh else: # For English and bilingual, use English only lighting_display = lighting lighting_zh = lighting objects_str = ', '.join([obj['class_name'] for obj in objects[:10]]) # CRITICAL: Emphasize brands EXTREMELY prominently - repeat multiple times if brands: brands_list = [b[0] for b in brands[:5]] brands_str = ', '.join(brands_list) brand_emphasis = f""" 🚨 CRITICAL BRAND REQUIREMENT 🚨 The following brands were POSITIVELY IDENTIFIED in this image: {brands_str} YOU ABSOLUTELY MUST: 1. Mention the brand name "{brands_list[0]}" explicitly in the FIRST sentence 2. Use the exact brand name - do not use generic terms like "bag" or "accessory" without the brand 3. Write naturally as if you're excited to share this {brands_list[0]} item 4. Example: "在傍晚光線下，這款{brands_list[0]}經典黑色菱格紋皮革包..." (CORRECT) 5. NOT acceptable: "在傍晚光線下，這款經典黑色菱格紋皮革包..." (WRONG - missing brand name!) THIS IS MANDATORY - The caption will be rejected if it doesn't mention {brands_str}. """ else: brands_str = 'None detected' brand_emphasis = "" # Enhanced scene description urban_scene = scene_info.get('urban', {}).get('top', 'unknown') mood = scene_info.get('mood', {}).get('top', 'neutral') comp_type = composition.get('composition_type', 'standard') context = f""" Analyze this image and generate an engaging, DETAILED social media caption with rich visual descriptions. **Visual Elements (Describe in Detail):** - Detected objects: {objects_str} - Scene composition: {comp_type} - Urban environment: {urban_scene} - **IMPORTANT**: Include specific details about: * Materials (leather, metal, fabric, canvas, etc.) * Colors (use descriptive terms: jet black, antique gold, midnight blue, etc.) * Textures (quilted, smooth, matte, glossy, metallic, etc.) * Design features (stitching patterns, hardware, logos, emblems, etc.) * Reflections and lighting effects on surfaces **Atmosphere:** - Lighting (analyzed with Places365 + CV): {lighting_display} (confidence: {lighting_confidence:.2f}) - Mood: {mood} **Brand Detection:** - Identified brands: {brands_str}{brand_emphasis} **Caption Structure (Required - BE SPECIFIC AND DETAILED):** 1. Opening hook - Most striking visual element with SPECIFIC details (1-2 sentences) {f"- 🚨 MANDATORY: Start with the BRAND NAME '{brands_list[0]}' in the FIRST sentence!" if brands else ""} {f"- Example (CORRECT): '這款{brands_list[0]}經典黑色菱格紋皮革包...'" if brands else ""} {f"- Example (WRONG): '這款經典黑色菱格紋皮革包...' (missing {brands_list[0]}!)" if brands else ""} - Be SPECIFIC: Include material, color, design features WITH the brand name 2. Visual details - Describe materials, textures, colors, and design elements (2-3 sentences) - Be SPECIFIC: mention quilting patterns, metal finishes, chain details, logo placements - Describe how light interacts with materials (reflections on leather, gleam of metal) - MUST use the EXACT lighting description: "{lighting_display}" 3. Atmospheric context - How lighting and mood create the scene's character (1-2 sentences) - Connect lighting to the overall visual impact - Describe depth, shadows, contrasts 4. Emotional connection & Engagement - How this resonates with viewers + call-to-action (1 sentence) **Content Requirements:** - Minimum information: 3-4 specific visual details per caption - Include material types, color descriptions, design features - Describe how lighting affects the appearance - Make it vivid and immersive Platform style: {platform_config['style']} """ # Language-specific examples with DETAILED visual descriptions AND BRAND NAMES if language == 'zh': brand_name_zh = brands_list[0] if brands else "Gucci" # Use detected brand or example example_correct = f"""正確範例 - 詳細描述 + 品牌提及 (繁體中文): "在{lighting_zh}的映襯下，這款{brand_name_zh}經典黑色菱格紋皮革包展現奢華質感，V字形縫線在柔軟小牛皮上勾勒出精緻的幾何圖案，復古金色雙G標誌在深色背景中熠熠生輝。金屬鏈條肩帶反射著{lighting_zh}，增添層次感與立體效果。皮革表面細膩的光澤與霧面質地形成迷人對比，每個細節都彰顯義大利工藝的極致追求。這樣的{brand_name_zh}單品不只是配件，更是品味與格調的完美詮釋。你的衣櫃裡有哪件經典單品？✨🖤" 注意：品牌名稱 "{brand_name_zh}" 出現在第一句！這是正確的做法。 CRITICAL: - 必須包含材質描述（皮革、金屬等） - 必須包含顏色細節（黑色、復古金色等） - 必須包含設計特點（縫線、標誌、鏈條等） - 必須使用"{lighting_zh}"來描述光線 """ elif language == 'en': brand_name_en = brands_list[0] if brands else "Gucci" # Use detected brand or example example_correct = f"""CORRECT EXAMPLE - Detailed Description + Brand Mention (ENGLISH ONLY - NO CHINESE): "Under the {lighting}, this {brand_name_en} classic black quilted leather bag showcases luxurious craftsmanship. V-shaped stitching traces intricate geometric patterns across supple calfskin, while the antique gold double-G logo gleams against the dark backdrop. The metal chain strap catches and reflects the {lighting}, adding dimension and depth to the piece. The leather surface presents a captivating contrast between fine sheen and matte texture, with every detail exemplifying Italian artisanship at its finest. This {brand_name_en} piece isn't just an accessory – it's a perfect expression of taste and sophistication. What's your timeless wardrobe essential? ✨🖤" NOTE: Brand name "{brand_name_en}" appears in the FIRST sentence! This is the correct approach. 🚨 ABSOLUTE REQUIREMENT FOR ENGLISH MODE 🚨 - Output must be 100% ENGLISH - zero Chinese characters allowed - MUST include material descriptions (leather, metal, etc.) - MUST include color details (black, antique gold, etc.) - MUST include design features (stitching, logo, chain, etc.) - MUST use "{lighting}" to describe the lighting - NO Chinese characters anywhere in the output """ else: # zh-en bilingual brand_name_en = brands_list[0] if brands else "Gucci" example_correct = f"""BILINGUAL EXAMPLE - 雙語範例: Caption in Traditional Chinese, with English hashtags support. (Details omitted for brevity) """ # Language-specific hashtag instructions if language == 'zh': hashtag_instruction = """ 【CRITICAL HASHTAG REQUIREMENT - 繁體中文】: - ALL hashtags MUST be in Traditional Chinese (繁體中文) - NEVER use English hashtags when language is 繁體中文 - Examples of CORRECT hashtags: ["時尚包包", "奢華風格", "皮革工藝", "精品配件"] - Examples of WRONG hashtags: ["FashionBlogger", "LuxuryLifestyle"] - DO NOT USE THESE """ elif language == 'en': hashtag_instruction = """ 【CRITICAL HASHTAG REQUIREMENT - English】: - ALL hashtags MUST be in English - NEVER use Chinese characters in hashtags - Examples of CORRECT hashtags: ["FashionBlogger", "LuxuryLifestyle", "LeatherCraft"] """ else: # zh-en hashtag_instruction = """ 【CRITICAL HASHTAG REQUIREMENT - Bilingual】: - Hashtags should MIX Traditional Chinese and English - First half in Chinese, second half in English - Example: ["時尚包包", "奢華風格", "FashionBlogger", "LuxuryLifestyle"] """ output_format = f""" Generate output in JSON format: {{ "caption": "string (minimum {platform_config['min_length']} chars, maximum {platform_config['max_length']} chars, engaging and descriptive)", "hashtags": ["tag1", "tag2", ...] ({platform_config['hashtag_count']} relevant hashtags), "tone": "casual|professional|playful", "platform": "{platform}" }} {hashtag_instruction} STRICT REQUIREMENTS: 1. Caption length: {platform_config['min_length']}-{platform_config['max_length']} characters 2. 🚨 EMOJI REQUIREMENT 🚨 - MUST use EXACTLY {platform_config['emoji_count']} emojis naturally integrated into caption text - Professional style: 1-2 emojis (e.g., ✨💼🌟) - Creative style: 2-3 emojis (e.g., 🎨✨💫🌙) - Authentic style: 2-3 emojis (e.g., 💖👜✨🖤) - Place emojis naturally within or at end of sentences 3. Caption must be pure descriptive text only - absolutely NO hashtags allowed 4. 🚨 CALL-TO-ACTION REQUIREMENT 🚨 - MUST include an engaging question or CTA at the end - Professional: Brief professional question (e.g., "What's your go-to piece?") - Creative: Thought-provoking question (e.g., "How does this speak to you?") - Authentic: Personal question (e.g., "What's your favorite timeless accessory?") 5. Write 3-4 complete sentences following the structure above 6. Be specific and vivid - describe what you see in detail 7. 【CRITICAL】 MUST use the EXACT lighting description: "{lighting_display}" - DO NOT substitute with similar terms - DO NOT use "金黃時刻" if the lighting is "{lighting_zh if language == 'zh' else lighting}" - DO NOT invent your own lighting description 8. 🚨 HASHTAG REQUIREMENT 🚨 - Generate {platform_config['hashtag_count']} relevant hashtags - Hashtags go ONLY in the 'hashtags' array, NEVER in the caption text - Mix of broad and specific tags - Include brand name as hashtag if detected 9. {"🚨 CRITICAL BRAND REQUIREMENT 🚨 - The brand name '" + brands_list[0] + "' MUST appear in the FIRST sentence of your caption. This is MANDATORY and NON-NEGOTIABLE. Example: " + ("'這款" + brands_list[0] + "經典黑色...'" if language == 'zh' else "'This " + brands_list[0] + " classic black...'") if brands else "No brands detected to mention"} 10. {"🚨 LANGUAGE REQUIREMENT 🚨 - Output must be 100% ENGLISH ONLY. NO Chinese characters allowed anywhere." if language == 'en' else ""} WRONG EXAMPLE (DO NOT DO THIS): "Lost in the city's towering skyscrapers 🏙️✨ | #UrbanVibes #CityLife" {example_correct} """ full_prompt = f"{system_instruction}\n\n{context}\n\n{output_format}" return full_prompt def generate_captions(self, analysis_results: Dict, image: Image.Image, platform: str = 'instagram', language: str = 'zh') -> List[Dict]: """Generate 3 captions with distinct styles: Professional, Creative, Authentic""" # Extract brands for style instructions brands_in_image = analysis_results.get('brands', []) brand_names = [b[0] for b in brands_in_image[:3]] if brands_in_image else [] brand_mention_requirement = f" CRITICAL: Mention {', '.join(brand_names)} brand(s) naturally in the caption." if brand_names else "" # Define 3 distinct styles styles = [ { 'name': 'professional', 'temp': 0.6, 'instruction': f'Professional style: Concise, elegant, sophisticated. Focus on quality and craftsmanship. Use refined language.{brand_mention_requirement}', 'length_modifier': 0.8 # Shorter, more concise }, { 'name': 'creative', 'temp': 0.7, 'instruction': f'Creative style: Artistic, expressive, imaginative. Use vivid metaphors and sensory descriptions. Balance detail with flair.{brand_mention_requirement}', 'length_modifier': 1.0 # Medium length }, { 'name': 'authentic', 'temp': 0.8, 'instruction': f'Authentic style: Personal, detailed, storytelling. Share rich observations and genuine feelings. Most descriptive and engaging.{brand_mention_requirement}', 'length_modifier': 1.2 # Longer, more detailed } ] variations = [] for style in styles: # Build style-specific prompt base_prompt = self.construct_prompt(analysis_results, platform, language) # Add style instruction style_prompt = f"""{base_prompt} **STYLE REQUIREMENT FOR THIS CAPTION:** {style['instruction']} Adjust tone to be clearly '{style['name']}' - this should be noticeably different from other styles.""" messages = [{ "role": "user", "content": [ {"type": "image", "image": image}, {"type": "text", "text": style_prompt} ] }] text = self.processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) image_inputs, video_inputs = process_vision_info(messages) inputs = self.processor( text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt" ) if torch.cuda.is_available(): inputs = inputs.to("cuda") # Generate with style-specific temperature config = self.generation_config.copy() config['temperature'] = style['temp'] with torch.no_grad(): generated_ids = self.model.generate(**inputs, **config) generated_ids_trimmed = [ out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids) ] output_text = self.processor.batch_decode( generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False )[0] parsed = self._parse_json_output(output_text) if parsed: # Force the correct tone parsed['tone'] = style['name'] # Remove any hashtags that leaked into caption if 'caption' in parsed: parsed['caption'] = self._remove_hashtags_from_caption(parsed['caption']) # Convert Simplified Chinese to Traditional if language is 'zh' if language == 'zh' or language == 'zh-en': parsed = self._convert_to_traditional(parsed) variations.append(parsed) return variations if variations else [self._get_fallback_caption(platform, language)] def _remove_hashtags_from_caption(self, caption: str) -> str: """Remove any hashtags, pipes, and debug info that leaked into caption text""" import re # CRITICAL FIX: Remove pipe symbol and everything after it (debug info) # Example: "Text 🕰️🌉 | SoftDiffusedLight" -> "Text 🕰️🌉" if '|' in caption: caption = caption.split('|')[0].strip() # Remove hashtags (words starting with #) caption = re.sub(r'#\w+', '', caption) caption = re.sub(r'#[\u4e00-\u9fff]+', '', caption) # Remove Chinese hashtags # Remove standalone weird text patterns (like "BLACKBELT") # If there's a suspicious all-caps word at the end without context, remove it words = caption.split() if len(words) > 0: last_word = words[-1].strip('✨💎👗🌟💫🖤') # If last "word" is all caps and doesn't look like a normal sentence word, remove it if last_word.isupper() and len(last_word) > 3 and not any(char in last_word for char in '.,!?'): caption = ' '.join(words[:-1]) # Remove excessive emojis at the end (more than 3) emoji_pattern = r'[\U0001F300-\U0001F9FF]{4,}$' caption = re.sub(emoji_pattern, '', caption) # Remove multiple spaces caption = re.sub(r'\s+', ' ', caption) # Remove trailing/leading whitespace caption = caption.strip() # Final cleanup: if caption ends with weird patterns like "✨X 👗💎", clean it if re.search(r'[✨💎👗🌟💫🖤]{2,}\s*$', caption): caption = re.sub(r'[✨💎👗🌟💫🖤\s]+$', '', caption).strip() return caption def _convert_to_traditional(self, caption: Dict) -> Dict: """Convert Simplified Chinese to Traditional Chinese""" if 'caption' in caption: caption['caption'] = self.cc.convert(caption['caption']) return caption def _parse_json_output(self, text: str) -> Dict: """Parse JSON output""" try: start = text.find('{') end = text.rfind('}') + 1 if start != -1 and end > start: json_str = text[start:end] return json.loads(json_str) except: pass return None def _get_fallback_caption(self, platform: str, language: str) -> Dict: """Fallback caption""" if language == 'en': return { 'caption': 'Every moment tells a story worth sharing. The world around us is filled with beauty waiting to be discovered. Take a pause and appreciate the details that make life extraordinary. What caught your eye today? ✨', 'hashtags': ['photography', 'daily', 'lifestyle', 'moment', 'capture'], 'tone': 'casual', 'platform': platform } else: return { 'caption': '每個瞬間都值得被記錄與分享。生活中充滿了等待被發現的美好細節。停下腳步，用心感受周遭的一切。今天什麼畫面觸動了你的心？✨', 'hashtags': ['攝影', '日常', '生活', '瞬間', '分享'], 'tone': 'casual', 'platform': platform } print("✓ CaptionGenerationManager (with Auto* classes for flexible model support) defined")