Update app.py

app.py

@@ -48,7 +48,7 @@ def cleanup_storage():
 
 TITLE = """
 <div style="text-align:center;margin:20px 0;">
-<h1>🎨 JoyCaption Three-Tone + Q&A (v3.2 ZeroGPU)</h1>
+<h1>🎨 JoyCaption Three-Tone + Q&A (v3.3 ZeroGPU)</h1>
 <p><strong>✅ Blog-friendly captions • SEO optimized • Longer descriptions • ZeroGPU optimized</strong></p>
 </div><hr>
 """
@@ -88,8 +88,8 @@ CAPTION_PROMPTS = {
         "max_tokens": 600
     },
     "uncensored_keywords": {
-        "system": "You are a detailed image captioner.",
-        "prompt": "Create a comprehensive description of this image. Describe the people, their appearance, clothing, poses, facial expressions, setting, and atmosphere in detail.",
+        "system": "You are a helpful image captioner.",
+        "prompt": "Describe this image.",
         "temperature": 0.6,
         "top_p": 0.9,
         "max_tokens": 600
@@ -137,9 +137,12 @@ def safe_generate_caption_direct(image, tone, max_chars=1000, keywords_text="",
         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 ONLY
-        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 and keywords for uncensored_keywords (third) tone ONLY
+        if tone == "uncensored_keywords":
+            if custom_instruction and custom_instruction.strip():
+                base_prompt += f" Make sure that you mention: {custom_instruction.strip()}"
+            if keywords_text and keywords_text.strip():
+                base_prompt += f" You are allowed to use these keywords: {keywords_text.strip()}"
         
         # Use conversation format
         convo = [
@@ -292,8 +295,8 @@ def generate_casual_friend_only(image, keywords_text="", custom_instruction=""):
 
 @spaces.GPU(duration=50)
 @torch.no_grad()
-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"
+def generate_uncensored_keywords_only(image, keywords_text="", custom_instruction=""):
+    result = safe_generate_caption_direct(image, "uncensored_keywords", max_chars=1000, keywords_text=keywords_text, custom_instruction=custom_instruction) if image else "❌ Upload image first"
     gc.collect()
     return result
 
@@ -341,6 +344,7 @@ def answer_question(image, question):
             temperature=0.3,  # Lower temperature for more accurate Q&A
             top_p=0.8,
             top_k=None,
+            repetition_penalty=1.1,
             use_cache=True,
             pad_token_id=processor.tokenizer.eos_token_id,
             eos_token_id=processor.tokenizer.eos_token_id
@@ -413,7 +417,7 @@ with gr.Blocks(title="JoyCaption Three-Tone + Q&A", theme=gr.themes.Soft()) as d
 
     g1.click(generate_engaging_only, [image_input, custom_instruction_input], out1)
     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)
+    g3.click(generate_uncensored_keywords_only, [image_input, keywords_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):