Update app.py

app.py

@@ -48,8 +48,8 @@ def cleanup_storage():
 
 TITLE = """
 <div style="text-align:center;margin:20px 0;">
-<h1>🎨 JoyCaption Three-Tone + Q&A (v2.7 ZeroGPU)</h1>
-<p><strong>✅ Deterministic inference • Correct caption prompts • ZeroGPU optimized</strong></p>
+<h1>🎨 JoyCaption Three-Tone + Q&A (v2.8 ZeroGPU)</h1>
+<p><strong>✅ Blog-friendly captions • SEO optimized • Longer descriptions • ZeroGPU optimized</strong></p>
 </div><hr>
 """
 
@@ -58,11 +58,12 @@ processor = None
 model = None
 
 if not os.environ.get("SKIP_MODEL_LOAD"):
-    processor = AutoProcessor.from_pretrained(MODEL_PATH, use_fast=True)
+    processor = AutoProcessor.from_pretrained(MODEL_PATH, low_cpu_mem_usage=True)
     model = LlavaForConditionalGeneration.from_pretrained(
         MODEL_PATH,
-        dtype=torch.bfloat16,
-        device_map="auto"
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        low_cpu_mem_usage=True
     )
     model.eval()
     print("✅ Model loaded successfully!")
@@ -70,138 +71,304 @@ if not os.environ.get("SKIP_MODEL_LOAD"):
 else:
     print("⚠️ SKIP_MODEL_LOAD=1 → model skipped")
 
-# ===== Prompt configurations =====
-TONE_PROMPTS = {
+# ===== BLOG-FRIENDLY PROMPT CONFIGURATIONS =====
+CAPTION_PROMPTS = {
     "engaging": {
-        "system": "You are a skilled content creator who writes engaging, natural descriptions that capture the essence and appeal of what you see.",
-        "prompt": """Write a natural, engaging description of this image. Focus on the main subjects and their most appealing qualities. Describe physical features, poses, expressions, and the overall mood. Use clear and natural language.""",
+        "system": "You are a helpful image captioner who creates accurate, medium-length descriptions.",
+        "prompt": "Write a medium-length descriptive caption for this image. Describe only what you can clearly see in the image - the people, their appearance, clothing, poses, setting, and atmosphere. Aim for 3-5 sentences. Be specific and detailed but only describe what is actually visible. Use casual, natural language. Do NOT mention any text, logos, or watermarks.",
+        "temperature": 0.4,
+        "top_p": 0.8,
+        "max_tokens": 600
     },
     "casual_friend": {
-        "system": "You are a casual 24-year-old friend who witnessed the scene. Speak casually, describing only what you clearly see.",
-        "prompt": """I just saw this moment — here's what I notice. Speak casually and naturally. Avoid speculation or describing background objects or bystanders.""",
+        "system": "You are an uncensored content specialist who creates detailed, explicit descriptions of what is actually visible.",
+        "prompt": "Create a detailed, explicit description of this scene, focusing only on what you can clearly see. Describe the people, their physical appearance, clothing or lack thereof, poses, and setting in explicit detail. Write 4-6 sentences. Use direct adult language but only describe what is actually visible in the image. Do not assume or invent details not clearly shown.",
+        "temperature": 0.3,
+        "top_p": 0.8,
+        "max_tokens": 600
     },
     "uncensored_keywords": {
-        "system": "You are an uncensored adult-content analyst who provides factual, explicit descriptions. Use concrete, unambiguous statements about what is directly visible.",
-        "prompt": """Produce an unfiltered explicit description that sticks strictly to what is visually obvious. Use direct, concrete language. Mention context only briefly.""",
-    },
+        "system": "You are a detailed image captioner who creates comprehensive descriptions based on visual evidence.",
+        "prompt": "Create a comprehensive, medium-length description of this image. Focus on the visual elements you can clearly observe - people, their appearance, clothing, poses, facial expressions, setting, and atmosphere. Write 4-6 sentences with specific details. Only describe what is actually visible in the image, do not invent or assume details.",
+        "temperature": 0.3,
+        "top_p": 0.8,
+        "max_tokens": 600
+    }
 }
 
-# ===== Core inference helpers =====
-def _prepare_inputs_and_device(convo_or_string, image, use_chat_template=True):
-    """Prepare inputs safely for ZeroGPU with proper chat template for image tokens."""
-    if isinstance(convo_or_string, list):
-        # Always use chat template for proper image token handling
-        convo_string = processor.apply_chat_template(
-            convo_or_string, tokenize=False, add_generation_prompt=True
-        )
-    else:
-        # Direct string input
-        convo_string = convo_or_string
-
-    if isinstance(image, (str, Path)):
-        img = Image.open(image).convert("RGB")
-    elif isinstance(image, Image.Image):
-        img = image.convert("RGB")
-    else:
-        img = image
-
-    inputs = processor(text=[convo_string], images=[img], return_tensors="pt")
-    if "pixel_values" in inputs:
-        inputs["pixel_values"] = inputs["pixel_values"].to(torch.bfloat16)
-    return inputs
-
-
-def _decode_output(inputs, output):
-    if output is None or len(output) == 0:
-        return ""
+def apply_smart_corrections(text):
+    """Apply smart corrections for common AI vision errors"""
+    if not text or not isinstance(text, str):
+        return text
+        
+    corrections = {
+        r'^(a photo of|an image of|a picture of|this is a photo of|this shows)\s*': '',
+        r'\btopless women\b': lambda m: 'nude women' if 'naked' in text.lower() or 'nude' in text.lower() else 'topless women',
+        r'\btopless woman\b': lambda m: 'nude woman' if 'naked' in text.lower() or 'nude' in text.lower() else 'topless woman',
+        r'\bthree women\b': lambda m: 'two women' if text.count('woman') + text.count('female') <= 2 else 'three women',
+        r'\bfour women\b': lambda m: 'three women' if text.count('woman') + text.count('female') <= 3 else 'four women',
+        r'\bwearing nothing\b': 'nude',
+        r'\bnot wearing.*clothes\b': 'nude',
+        r'\bcompletely naked\b': 'nude',
+        r'\bfully nude\b': 'nude',
+    }
+    
+    corrected_text = text
     try:
-        input_len = inputs.get("input_ids", torch.empty((1, 0))).shape[1]
-        decoded = processor.tokenizer.decode(
-            output[0][input_len:], skip_special_tokens=True, clean_up_tokenization_spaces=True
-        )
-        return decoded.strip()
-    except Exception:
-        try:
-            return processor.tokenizer.decode(output[0], skip_special_tokens=True).strip()
-        except:
-            return ""
-
-
-def cleanup_after_inference():
-    gc.collect()
-
+        for pattern, replacement in corrections.items():
+            if callable(replacement):
+                corrected_text = re.sub(pattern, replacement, corrected_text, flags=re.IGNORECASE)
+            else:
+                corrected_text = re.sub(pattern, replacement, corrected_text, flags=re.IGNORECASE)
+    except Exception as e:
+        print(f"Error in corrections: {e}")
+        return text
+    
+    return corrected_text
 
-def run_image_chat_generation(convo, image, max_new_tokens=150, caption_mode=False):
-    """Deterministic generation with correct mode separation."""
-    if processor is None or model is None:
-        return None, "❌ Model not initialized"
+def safe_generate_caption_direct(image, tone, max_chars=1000, keywords_text="", custom_instruction=""):
+    """Generate caption using blog-friendly approach"""
     try:
-        inputs = _prepare_inputs_and_device(convo, image, use_chat_template=True)
+        if image is None:
+            return f"❌ No image provided for {tone}"
+        
+        caption_config = CAPTION_PROMPTS.get(tone, CAPTION_PROMPTS["engaging"])
+        base_prompt = caption_config["prompt"]
+        
+        # Handle keywords for casual_friend (erotic) tone
+        if tone == "casual_friend" and keywords_text and keywords_text.strip():
+            base_prompt += f" Pay special attention to these elements if present: {keywords_text.strip()}"
+        
+        # Handle custom instructions for uncensored_keywords (third) tone
+        if tone == "uncensored_keywords" and custom_instruction and custom_instruction.strip():
+            base_prompt += f" Make sure that you mention: {custom_instruction.strip()}"
+        
+        # Handle custom instructions for other tones
+        if tone != "uncensored_keywords" and custom_instruction and custom_instruction.strip():
+            base_prompt += f" Also focus on: {custom_instruction.strip()}"
+        
+        # Use conversation format
+        convo = [
+            {"role": "system", "content": caption_config["system"]},
+            {"role": "user", "content": base_prompt}
+        ]
+        
+        convo_string = processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
+        inputs = processor(text=[convo_string], images=[image], return_tensors="pt")
+        
+        device = next(model.parameters()).device
+        inputs = {k: v.to(device, non_blocking=True) if hasattr(v, 'to') else v for k, v in inputs.items()}
+        
+        if 'pixel_values' in inputs:
+            inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16)
+        
+        temperature = caption_config.get("temperature", 0.4)
+        top_p = caption_config.get("top_p", 0.8)
+        max_tokens = caption_config.get("max_tokens", 600)
+        
         with torch.no_grad():
             output = model.generate(
                 **inputs,
-                max_new_tokens=max_new_tokens,
-                do_sample=False,
-                temperature=None,
-                top_p=None,
-                repetition_penalty=1.0,
-                pad_token_id=processor.tokenizer.eos_token_id,
-                eos_token_id=processor.tokenizer.eos_token_id,
+                max_new_tokens=max_tokens,
+                do_sample=True,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=None,
                 use_cache=True,
+                pad_token_id=processor.tokenizer.eos_token_id,
+                eos_token_id=processor.tokenizer.eos_token_id
             )
-        decoded = _decode_output(inputs, output)
-        cleanup_after_inference()
-        return decoded, None
+        
+        if output is None or len(output) == 0:
+            return f"❌ No output generated for {tone}"
+        
+        # Proper decoding
+        if 'input_ids' in inputs and len(inputs['input_ids'].shape) >= 2:
+            input_length = inputs['input_ids'].shape[1]
+            if len(output[0]) > input_length:
+                generate_ids = output[0][input_length:]
+                result = processor.tokenizer.decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+            else:
+                result = processor.tokenizer.decode(output[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        else:
+            result = processor.tokenizer.decode(output[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        
+        result = result.strip()
+        result = apply_smart_corrections(result)
+        
+        # Cleanup after generation
+        del inputs, output
+        gc.collect()
+        
+        # Apply postprocessing
+        final_result = postprocess_caption(result, max_chars=max_chars)
+        
+        return final_result if final_result else f"❌ Empty result for {tone}"
+        
     except Exception as e:
-        cleanup_after_inference()
-        return None, f"❌ Generation error: {str(e)[:200]}"
+        gc.collect()
+        return f"❌ Error: {str(e)[:200]}"
 
-# ===== Caption generators =====
-def safe_generate_caption_direct(image, tone, max_chars=600, keywords_text="", custom_instruction=""):
-    if image is None:
-        return f"❌ No image provided for {tone}"
-    tone_cfg = TONE_PROMPTS.get(tone, TONE_PROMPTS["engaging"])
-    
-    # Build proper conversation format for captioning (JoyCaption expects this format)
-    base_prompt = tone_cfg["prompt"]
-    if tone=="uncensored_keywords" and keywords_text.strip():
-        base_prompt += f"\n\nKeywords to mention: {keywords_text.strip()}"
-    if custom_instruction.strip():
-        base_prompt += f"\n\nInclude: {custom_instruction.strip()}"
-    
-    # Use conversation format as per HuggingFace example
-    convo = [
-        {"role": "system", "content": tone_cfg["system"]},
-        {"role": "user", "content": base_prompt}
-    ]
-    decoded, err = run_image_chat_generation(convo, image, caption_mode=True)
-    if err: return err
-    result = postprocess_caption(decoded, max_chars=max_chars)
-    return result if result else "❌ Empty result"
+def safe_generate_custom_prompt(image, system_prompt, user_prompt, max_chars=1000):
+    """Generate caption using custom system and user prompts for playground"""
+    try:
+        if image is None:
+            return "❌ No image provided"
+        
+        if not system_prompt or not system_prompt.strip():
+            return "❌ System prompt is required"
+            
+        if not user_prompt or not user_prompt.strip():
+            return "❌ User prompt is required"
+        
+        # Use custom prompts
+        convo = [
+            {"role": "system", "content": system_prompt.strip()},
+            {"role": "user", "content": user_prompt.strip()}
+        ]
+        
+        convo_string = processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
+        inputs = processor(text=[convo_string], images=[image], return_tensors="pt")
+        
+        device = next(model.parameters()).device
+        inputs = {k: v.to(device, non_blocking=True) if hasattr(v, 'to') else v for k, v in inputs.items()}
+        
+        if 'pixel_values' in inputs:
+            inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16)
+        
+        with torch.no_grad():
+            output = model.generate(
+                **inputs,
+                max_new_tokens=600,
+                do_sample=True,
+                temperature=0.5,
+                top_p=0.8,
+                top_k=None,
+                use_cache=True,
+                pad_token_id=processor.tokenizer.eos_token_id,
+                eos_token_id=processor.tokenizer.eos_token_id
+            )
+        
+        if output is None or len(output) == 0:
+            return "❌ No output generated"
+        
+        # Proper decoding
+        if 'input_ids' in inputs and len(inputs['input_ids'].shape) >= 2:
+            input_length = inputs['input_ids'].shape[1]
+            if len(output[0]) > input_length:
+                generate_ids = output[0][input_length:]
+                result = processor.tokenizer.decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+            else:
+                result = processor.tokenizer.decode(output[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        else:
+            result = processor.tokenizer.decode(output[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        
+        result = result.strip()
+        result = apply_smart_corrections(result)
+        
+        # Cleanup after generation
+        del inputs, output
+        gc.collect()
+        
+        # Apply postprocessing
+        final_result = postprocess_caption(result, max_chars=max_chars)
+        
+        return final_result if final_result else "❌ Empty result"
+        
+    except Exception as e:
+        gc.collect()
+        return f"❌ Error: {str(e)[:200]}"
 
+# Individual GPU-decorated functions for CAPTIONS
+@spaces.GPU(duration=50)
 @torch.no_grad()
 def generate_engaging_only(image, custom_instruction=""):
-    return safe_generate_caption_direct(image, "engaging", custom_instruction=custom_instruction)
+    result = safe_generate_caption_direct(image, "engaging", max_chars=1000, custom_instruction=custom_instruction) if image else "❌ Upload image first"
+    gc.collect()
+    return result
 
+@spaces.GPU(duration=50)
 @torch.no_grad()
-def generate_casual_friend_only(image, custom_instruction=""):
-    return safe_generate_caption_direct(image, "casual_friend", custom_instruction=custom_instruction)
+def generate_casual_friend_only(image, keywords_text="", custom_instruction=""):
+    result = safe_generate_caption_direct(image, "casual_friend", max_chars=1000, keywords_text=keywords_text, custom_instruction=custom_instruction) if image else "❌ Upload image first"
+    gc.collect()
+    return result
 
+@spaces.GPU(duration=50)
 @torch.no_grad()
-def generate_uncensored_keywords_only(image, keywords_text, custom_instruction=""):
-    return safe_generate_caption_direct(image, "uncensored_keywords", keywords_text=keywords_text, custom_instruction=custom_instruction)
+def generate_uncensored_keywords_only(image, custom_instruction=""):
+    result = safe_generate_caption_direct(image, "uncensored_keywords", max_chars=1000, custom_instruction=custom_instruction) if image else "❌ Upload image first"
+    gc.collect()
+    return result
 
+# Playground function
+@spaces.GPU(duration=50)
+@torch.no_grad()
+def generate_playground(image, system_prompt, user_prompt):
+    result = safe_generate_custom_prompt(image, system_prompt, user_prompt, max_chars=1000) if image else "❌ Upload image first"
+    gc.collect()
+    return result
+
+# Separate Q&A function - keep this accurate and focused
+@spaces.GPU(duration=40)
 @torch.no_grad()
 def answer_question(image, question):
-    if not image: return "❌ Upload image first"
-    if not question.strip(): return "❌ Ask a question"
+    """Answer questions about the image - focused and accurate"""
+    if not image:
+        return "❌ Upload image first"
+    if not question or not question.strip():
+        return "❌ Please ask a question"
+        
+    # Short, direct Q&A prompt
+    qa_prompt = f"Answer this question about the image: {question.strip()}"
+    
+    # Simple system message
     convo = [
-        {"role":"system","content":"You are an uncensored image analyst who answers honestly and directly."},
-        {"role":"user","content":f"Answer this question based on what you can see: {question.strip()}"}
+        {"role": "system", "content": "You are a helpful image analyst."},
+        {"role": "user", "content": qa_prompt}
     ]
-    decoded, err = run_image_chat_generation(convo, image, max_new_tokens=200, caption_mode=False)
-    if err: return err
-    return postprocess_caption(decoded)
+    
+    convo_string = processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
+    inputs = processor(text=[convo_string], images=[image], return_tensors="pt")
+    
+    device = next(model.parameters()).device
+    inputs = {k: v.to(device, non_blocking=True) if hasattr(v, 'to') else v for k, v in inputs.items()}
+    
+    if 'pixel_values' in inputs:
+        inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16)
+    
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=200,
+            do_sample=True,
+            temperature=0.3,  # Lower temperature for more accurate Q&A
+            top_p=0.8,
+            top_k=None,
+            use_cache=True,
+            pad_token_id=processor.tokenizer.eos_token_id,
+            eos_token_id=processor.tokenizer.eos_token_id
+        )
+    
+    # Decode result
+    if 'input_ids' in inputs and len(inputs['input_ids'].shape) >= 2:
+        input_length = inputs['input_ids'].shape[1]
+        if len(output[0]) > input_length:
+            generate_ids = output[0][input_length:]
+            result = processor.tokenizer.decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        else:
+            result = processor.tokenizer.decode(output[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    else:
+        result = processor.tokenizer.decode(output[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    
+    result = result.strip()
+    
+    # Cleanup
+    del inputs, output
+    gc.collect()
+    
+    final_result = postprocess_caption(result, max_chars=300)
+    return final_result if final_result else "❌ No answer generated"
 
 # ===== Export =====
 def export_joycaption_data(keywords, custom_instructions, question, engaging_caption, casual_caption, keywords_caption, qa_answer, image_path=""):
@@ -231,26 +398,26 @@ with gr.Blocks(title="JoyCaption Three-Tone + Q&A", theme=gr.themes.Soft()) as d
     gr.HTML(TITLE)
     with gr.Row():
         with gr.Column(scale=1):
-            image_input = gr.Image(type="filepath", label="📸 Upload Image", height=400)
+            image_input = gr.Image(type="pil", label="📸 Upload Image", height=400)
             keywords_input = gr.Textbox(label="🏷️ Keywords", lines=2)
             custom_instruction_input = gr.Textbox(label="🎯 Custom Instruction", lines=2)
             question_input = gr.Textbox(label="❓ Ask Question", lines=2)
             ask_btn = gr.Button("❓ Ask", variant="secondary")
-            qa_output = gr.Textbox(label="Q&A", lines=5, show_copy_button=True)
+            qa_output = gr.Textbox(label="Q&A", lines=4, show_copy_button=True)
         with gr.Column(scale=1):
-            g1 = gr.Button("✨ Engaging", variant="primary")
-            out1 = gr.Textbox(lines=5, show_copy_button=True)
-            g2 = gr.Button("😎 Casual Friend", variant="primary")
-            out2 = gr.Textbox(lines=5, show_copy_button=True)
-            g3 = gr.Button("🔴 Keywords", variant="secondary")
-            out3 = gr.Textbox(lines=5, show_copy_button=True)
+            g1 = gr.Button("📝 Casual Descriptive", variant="primary")
+            out1 = gr.Textbox(lines=7, show_copy_button=True)
+            g2 = gr.Button("🔥 Erotic", variant="secondary")
+            out2 = gr.Textbox(lines=7, show_copy_button=True)
+            g3 = gr.Button("🎯 Custom Instruction", variant="secondary")
+            out3 = gr.Textbox(lines=7, show_copy_button=True)
             export_btn = gr.Button("📥 Export All Data")
             export_out = gr.Textbox(visible=False)
             export_file = gr.File(visible=False)
 
     g1.click(generate_engaging_only, [image_input, custom_instruction_input], out1)
-    g2.click(generate_casual_friend_only, [image_input, custom_instruction_input], out2)
-    g3.click(generate_uncensored_keywords_only, [image_input, keywords_input, custom_instruction_input], out3)
+    g2.click(generate_casual_friend_only, [image_input, keywords_input, custom_instruction_input], out2)
+    g3.click(generate_uncensored_keywords_only, [image_input, custom_instruction_input], out3)
     ask_btn.click(answer_question, [image_input, question_input], qa_output)
 
     def handle_export(k, c, q, e1, e2, e3, qa, img):
@@ -264,5 +431,34 @@ with gr.Blocks(title="JoyCaption Three-Tone + Q&A", theme=gr.themes.Soft()) as d
 
     export_btn.click(handle_export, [keywords_input, custom_instruction_input, question_input, out1, out2, out3, qa_output, image_input], [export_out, export_file])
 
+    # ===== PLAYGROUND SECTION =====
+    gr.HTML("<hr><h2>🧪 Playground - Custom Prompts</h2><p>Test custom system and user prompts (not included in JSON export)</p>")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            playground_system = gr.Textbox(
+                label="🔧 System Prompt",
+                lines=2,
+                value="You are a helpful image captioner who creates accurate, detailed descriptions based only on what is clearly visible.",
+                placeholder="Enter custom system prompt..."
+            )
+            playground_prompt = gr.Textbox(
+                label="💬 User Prompt", 
+                lines=3,
+                value="Describe what you can clearly see in this image. Focus on the people present, their hair colors, body types, facial expressions, clothing or lack thereof, poses, and actions. Write 4-6 sentences with specific details. Only describe what is actually visible - do not assume or invent details not clearly shown.",
+                placeholder="Enter custom user prompt..."
+            )
+            playground_btn = gr.Button("🧪 Generate", variant="secondary")
+        
+        with gr.Column(scale=1):
+            playground_output = gr.Textbox(
+                label="🎯 Playground Output",
+                lines=7,
+                show_copy_button=True,
+                placeholder="Custom prompt results will appear here..."
+            )
+
+    playground_btn.click(generate_playground, [image_input, playground_system, playground_prompt], playground_output)
+
 if __name__ == "__main__":
     demo.launch()