Update app.py

app.py

@@ -16,7 +16,6 @@ except Exception:
 import gradio as gr
 import torch
 from transformers import LlavaForConditionalGeneration, AutoProcessor
-from PIL import Image
 import tempfile, gc, os, shutil, json, time, re
 from urllib.parse import urlparse
 from typing import Optional
@@ -42,9 +41,7 @@ def fix_image_url(raw_url_or_path: str, host: Optional[str] = None) -> str:
         host = host.rstrip("/")
         if not (host.startswith("http://") or host.startswith("https://")):
             host = "https://" + host
-        p = raw_url_or_path
-        if p.startswith("/"):
-            p = p[1:]
+        p = raw_url_or_path.lstrip("/")
         return f"{host}/gradio_api/file=/{p}"
     return raw_url_or_path
 
@@ -69,49 +66,34 @@ def force_clear_all_caches():
             torch.cuda.empty_cache()
             torch.cuda.synchronize()
         gc.collect()
-    except Exception as e:
-        print(f"⚠️ Cache clear warning: {e}")
+    except Exception:
+        pass
 
 force_clear_all_caches()
 
 # ===== SETUP =====
 _tmpdir = tempfile.gettempdir()
-os.environ["HF_HOME"] = os.path.join(_tmpdir, "hf_cache")
-os.environ["TRANSFORMERS_CACHE"] = os.path.join(_tmpdir, "transformers_cache") 
-os.environ["HF_DATASETS_CACHE"] = os.path.join(_tmpdir, "datasets_cache")
-os.environ["TORCH_HOME"] = os.path.join(_tmpdir, "torch_cache")
+for key, folder in {
+    "HF_HOME": "hf_cache",
+    "TRANSFORMERS_CACHE": "transformers_cache",
+    "HF_DATASETS_CACHE": "datasets_cache",
+    "TORCH_HOME": "torch_cache"
+}.items():
+    os.environ[key] = os.path.join(_tmpdir, folder)
 
 MODEL_PATH = "fancyfeast/llama-joycaption-beta-one-hf-llava"
 SPACE_HOST = os.environ.get("SPACE_HOST") or os.environ.get("HF_SPACE_HOST") or None
 
-def cleanup_storage():
-    try:
-        for key in ["HF_HOME","TRANSFORMERS_CACHE","HF_DATASETS_CACHE","TORCH_HOME"]:
-            d = os.environ.get(key)
-            if d and os.path.exists(d):
-                shutil.rmtree(d, ignore_errors=True)
-        gc.collect()
-    except Exception as e:
-        print(f"⚠️ Storage cleanup warning: {e}")
-
-TITLE = """
-<div style="text-align:center;margin:20px 0;">
-<h1>🎨 JoyCaption Advanced Prompting System (v6.0)</h1>
-<p><strong>🎛️ Fully customizable prompts • Template helpers • Professional control</strong></p>
-</div><hr>
-"""
-
 print("🚀 Loading JoyCaption model...")
 processor = AutoProcessor.from_pretrained(MODEL_PATH, cache_dir=None)
 model = LlavaForConditionalGeneration.from_pretrained(
-    MODEL_PATH, 
-    torch_dtype=torch.bfloat16, 
+    MODEL_PATH,
+    torch_dtype=torch.bfloat16,
     device_map="auto",
     cache_dir=None,
     low_cpu_mem_usage=True
 )
 model.eval()
-cleanup_storage()
 print("✅ Model loaded successfully!")
 
 # ===== DEFAULT PROMPTS =====
@@ -130,6 +112,7 @@ DEFAULT_PROMPTS = {
     }
 }
 
+# ===== CORE CAPTIONING =====
 def safe_generate_caption_direct(image, system_prompt, user_prompt, max_chars=1200):
     try:
         if image is None:
@@ -144,7 +127,7 @@ def safe_generate_caption_direct(image, system_prompt, user_prompt, max_chars=12
         ]
         convo_string = processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
         inputs = processor(text=[convo_string], images=[image], return_tensors="pt").to("cuda")
-        inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16)
+        inputs["pixel_values"] = inputs["pixel_values"].to(torch.bfloat16)
         with torch.no_grad():
             output = model.generate(
                 **inputs,
@@ -152,14 +135,11 @@ def safe_generate_caption_direct(image, system_prompt, user_prompt, max_chars=12
                 do_sample=True,
                 temperature=0.6,
                 top_p=0.9,
-                top_k=None,
                 use_cache=True,
                 pad_token_id=processor.tokenizer.eos_token_id,
-                eos_token_id=processor.tokenizer.eos_token_id
+                eos_token_id=processor.tokenizer.eos_token_id,
             )
-        if not output or len(output) == 0:
-            return "❌ No output generated"
-        input_length = inputs['input_ids'].shape[1]
+        input_length = inputs["input_ids"].shape[1]
         result = processor.tokenizer.decode(output[0][input_length:], skip_special_tokens=True)
         del inputs, output
         torch.cuda.empty_cache()
@@ -170,40 +150,34 @@ def safe_generate_caption_direct(image, system_prompt, user_prompt, max_chars=12
         gc.collect()
         return f"❌ Error: {str(e)[:200]}"
 
-# ===== CAPTION FUNCTIONS =====
-@spaces.GPU(duration=60)
-@torch.no_grad()
-def generate_caption_1(image, system1, user1):
-    if not image: return "❌ Upload image first"
-    return safe_generate_caption_direct(image, system1, user1)
-
+# ===== INDIVIDUAL CAPTION WRAPPERS =====
 @spaces.GPU(duration=60)
 @torch.no_grad()
-def generate_caption_2(image, system2, user2):
-    if not image: return "❌ Upload image first"
-    return safe_generate_caption_direct(image, system2, user2)
-
-@spaces.GPU(duration=60)
-@torch.no_grad()
-def generate_caption_3(image, system3, user3):
-    if not image: return "❌ Upload image first"
-    return safe_generate_caption_direct(image, system3, user3)
+def generate_caption(image, system, user):
+    if not image:
+        return "❌ Upload image first"
+    return safe_generate_caption_direct(image, system, user)
 
+# ===== Q&A =====
 @spaces.GPU(duration=40)
 @torch.no_grad()
 def answer_question(image, question):
-    if not image: return "❌ Upload image first"
-    if not question.strip(): return "❌ Please ask a question"
+    if not image:
+        return "❌ Upload image first"
+    if not question.strip():
+        return "❌ Please ask a question"
     try:
         torch.cuda.empty_cache()
         gc.collect()
-        convo = [{"role": "system", "content": "You are a helpful image captioner."},
-                 {"role": "user", "content": question.strip()}]
+        convo = [
+            {"role": "system", "content": "You are a helpful image captioner."},
+            {"role": "user", "content": question.strip()},
+        ]
         convo_string = processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
         inputs = processor(text=[convo_string], images=[image], return_tensors="pt").to("cuda")
-        inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16)
+        inputs["pixel_values"] = inputs["pixel_values"].to(torch.bfloat16)
         output = model.generate(**inputs, max_new_tokens=300, do_sample=True, temperature=0.6, top_p=0.9)
-        input_length = inputs['input_ids'].shape[1]
+        input_length = inputs["input_ids"].shape[1]
         result = processor.tokenizer.decode(output[0][input_length:], skip_special_tokens=True)
         del inputs, output
         torch.cuda.empty_cache()
@@ -214,72 +188,81 @@ def answer_question(image, question):
         gc.collect()
         return f"❌ Q&A Error: {str(e)[:200]}"
 
-# ===== TEMPLATE FUNCTIONS =====
+# ===== TEMPLATE HELPERS =====
 def insert_template(current_text, template_text, field_content):
-    if not field_content.strip(): return current_text
-    formatted_template = template_text.format(content=field_content.strip())
-    if formatted_template in current_text: return current_text
-    return (current_text.rstrip() + " " + formatted_template).strip()
+    if not field_content.strip():
+        return current_text
+    formatted = template_text.format(content=field_content.strip())
+    if formatted in current_text:
+        return current_text
+    return (current_text.rstrip() + " " + formatted).strip()
 
 def create_template_functions():
-    def insert_key(system_text, user_text, keywords_content):
-        template = "Pay attention to these keywords: {content}."
-        return (insert_template(system_text, template, keywords_content),
-                insert_template(user_text, template, keywords_content))
-    def insert_que(system_text, user_text, question_content):
-        template = "Answer this question: {content}."
-        return (insert_template(system_text, template, question_content),
-                insert_template(user_text, template, question_content))
-    def insert_use(system_text, user_text, custom_content):
-        template = "Make sure that you mention: {content}."
-        return (insert_template(system_text, template, custom_content),
-                insert_template(user_text, template, custom_content))
-    def insert_not(system_text, user_text, avoid_content):
-        template = "Do NOT mention: {content}."
-        return (insert_template(system_text, template, avoid_content),
-                insert_template(user_text, template, avoid_content))
+    def insert_key(s, u, c):
+        t = "Pay attention to these keywords: {content}."
+        return insert_template(s, t, c), insert_template(u, t, c)
+    def insert_que(s, u, c):
+        t = "Answer this question: {content}."
+        return insert_template(s, t, c), insert_template(u, t, c)
+    def insert_use(s, u, c):
+        t = "Make sure that you mention: {content}."
+        return insert_template(s, t, c), insert_template(u, t, c)
+    def insert_not(s, u, c):
+        t = "Do NOT mention: {content}."
+        return insert_template(s, t, c), insert_template(u, t, c)
     return insert_key, insert_que, insert_use, insert_not
 
 # ===== EXPORT =====
-def export_joycaption_data(keywords, custom_instructions, avoid, question, cap1, cap2, cap3, qa_answer, image_path=""):
+def export_joycaption_data(keywords, custom_instructions, avoid, question, c1, c2, c3, qa, img):
     try:
-        data = {"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), "source":"JoyCaption","data":{}}
-        if keywords.strip(): data["data"]["keywords"]=keywords.strip()
-        if custom_instructions.strip(): data["data"]["custom_instructions"]=custom_instructions.strip()
-        if avoid.strip(): data["data"]["avoid"]=avoid.strip()
-        if question.strip(): data["data"]["question"]=question.strip()
-        if image_path.strip():
-            data["data"]["image_local_path"]=image_path
-            image_url=fix_image_url(image_path, host=(SPACE_HOST or ""))
-            if image_url: data["data"]["image_url"]=image_url
-        if cap1.strip(): data["data"]["caption_casual"]=cap1.strip()
-        if cap2.strip(): data["data"]["caption_friendly"]=cap2.strip()
-        if cap3.strip(): data["data"]["caption_erotic"]=cap3.strip()
-        if qa_answer.strip(): data["data"]["qa_answer"]=qa_answer.strip()
-        if not data["data"]: return "❌ No data to export", None
+        data = {"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), "source": "JoyCaption", "data": {}}
+        add = data["data"]
+        if keywords.strip(): add["keywords"] = keywords.strip()
+        if custom_instructions.strip(): add["custom_instructions"] = custom_instructions.strip()
+        if avoid.strip(): add["avoid"] = avoid.strip()
+        if question.strip(): add["question"] = question.strip()
+        if img.strip():
+            add["image_local_path"] = img
+            url = fix_image_url(img, host=(SPACE_HOST or ""))
+            if url: add["image_url"] = url
+        if c1.strip(): add["caption_casual"] = c1.strip()
+        if c2.strip(): add["caption_friendly"] = c2.strip()
+        if c3.strip(): add["caption_erotic"] = c3.strip()
+        if qa.strip(): add["qa_answer"] = qa.strip()
+        if not add:
+            return "❌ No data to export", None
         js = json.dumps(data, indent=2, ensure_ascii=False)
         fn = f"joycaption_{time.strftime('%Y%m%d_%H%M%S')}.json"
-        return f"✅ Exported {len(data['data'])} fields", (js, fn)
+        return f"✅ Exported {len(add)} fields", (js, fn)
     except Exception as e:
-        return f"❌ Export failed: {str(e)}", None
+        return f"❌ Export failed: {e}", None
 
 # ===== UI =====
 with gr.Blocks(title="JoyCaption Advanced Prompting System", theme=gr.themes.Soft()) as demo:
-    gr.HTML(TITLE)
+    gr.HTML("""
+    <style>
+    textarea.autoresize {overflow-y:hidden!important;min-height:50px!important;height:auto!important;}
+    </style>
+    """)
+    gr.HTML("<h1 style='text-align:center;margin-top:10px;'>🎨 JoyCaption Advanced Prompting System (v6.0)</h1><hr>")
+
     insert_key, insert_que, insert_use, insert_not = create_template_functions()
+    active_tab = gr.State("casual")
 
     with gr.Row():
+        # LEFT
         with gr.Column(scale=1):
             image_input = gr.Image(type="pil", label="📸 Image", height=400)
             keywords_input = gr.Textbox(label="🏷️ Keywords", lines=2, placeholder="e.g. beach, sunset")
-            custom_instruction_input = gr.Textbox(label="🎯 Custom", lines=2, placeholder="Add extra instructions")
-            avoid_input = gr.Textbox(label="🚫 Avoid", lines=2, placeholder="What to avoid")
+            custom_instruction_input = gr.Textbox(label="🎯 Custom", lines=2, placeholder="Extra instructions")
+            avoid_input = gr.Textbox(label="🚫 Avoid", lines=2, placeholder="Things to avoid")
             question_input = gr.Textbox(label="❓ Question", lines=2, placeholder="Ask about image")
             ask_btn = gr.Button("Ask", variant="secondary")
             qa_output = gr.Textbox(label="Answer", lines=3, show_copy_button=True)
 
+        # RIGHT
         with gr.Column(scale=1):
-            # Template buttons moved above tabs
+            # template buttons top
             gr.Markdown("**Insert Template**")
             with gr.Row():
                 key_btn = gr.Button("key", size="sm")
@@ -287,22 +270,26 @@ with gr.Blocks(title="JoyCaption Advanced Prompting System", theme=gr.themes.Sof
                 use_btn = gr.Button("use", size="sm")
                 not_btn = gr.Button("not", size="sm")
 
-            # Tabs with flexible Textboxes
-            with gr.Tab("📝 Casual"):
-                system1 = gr.Textbox(label="System", show_label=True, placeholder="", value=DEFAULT_PROMPTS["casual"]["system"], lines=1, max_lines=5)
-                user1 = gr.Textbox(label="User", show_label=True, placeholder="", value=DEFAULT_PROMPTS["casual"]["user"], lines=1, max_lines=8)
+            # Tabs
+            with gr.Tab("📝 Casual") as tab1:
+                gr.Markdown("**System Prompt**")
+                system1 = gr.Textbox(show_label=False, value=DEFAULT_PROMPTS["casual"]["system"], elem_classes="autoresize")
+                gr.Markdown("**User Prompt**")
+                user1 = gr.Textbox(show_label=False, value=DEFAULT_PROMPTS["casual"]["user"], elem_classes="autoresize")
                 gen1_btn = gr.Button("Generate Casual", variant="primary")
                 out1 = gr.Textbox(lines=5, show_copy_button=True)
-
-            with gr.Tab("🤝 Friendly"):
-                system2 = gr.Textbox(label="System", show_label=True, placeholder="", value=DEFAULT_PROMPTS["friendly"]["system"], lines=1, max_lines=5)
-                user2 = gr.Textbox(label="User", show_label=True, placeholder="", value=DEFAULT_PROMPTS["friendly"]["user"], lines=1, max_lines=8)
+            with gr.Tab("🤝 Friendly") as tab2:
+                gr.Markdown("**System Prompt**")
+                system2 = gr.Textbox(show_label=False, value=DEFAULT_PROMPTS["friendly"]["system"], elem_classes="autoresize")
+                gr.Markdown("**User Prompt**")
+                user2 = gr.Textbox(show_label=False, value=DEFAULT_PROMPTS["friendly"]["user"], elem_classes="autoresize")
                 gen2_btn = gr.Button("Generate Friendly", variant="primary")
                 out2 = gr.Textbox(lines=5, show_copy_button=True)
-
-            with gr.Tab("🔥 Erotic"):
-                system3 = gr.Textbox(label="System", show_label=True, placeholder="", value=DEFAULT_PROMPTS["erotic"]["system"], lines=1, max_lines=5)
-                user3 = gr.Textbox(label="User", show_label=True, placeholder="", value=DEFAULT_PROMPTS["erotic"]["user"], lines=1, max_lines=8)
+            with gr.Tab("🔥 Erotic") as tab3:
+                gr.Markdown("**System Prompt**")
+                system3 = gr.Textbox(show_label=False, value=DEFAULT_PROMPTS["erotic"]["system"], elem_classes="autoresize")
+                gr.Markdown("**User Prompt**")
+                user3 = gr.Textbox(show_label=False, value=DEFAULT_PROMPTS["erotic"]["user"], elem_classes="autoresize")
                 gen3_btn = gr.Button("Generate Erotic", variant="primary")
                 out3 = gr.Textbox(lines=5, show_copy_button=True)
 
@@ -311,27 +298,54 @@ with gr.Blocks(title="JoyCaption Advanced Prompting System", theme=gr.themes.Sof
             export_out = gr.Textbox(visible=False)
             export_file = gr.File(visible=False)
 
+    # === TAB SWITCH HANDLERS ===
+    tab1.select(lambda: "casual", None, active_tab)
+    tab2.select(lambda: "friendly", None, active_tab)
+    tab3.select(lambda: "erotic", None, active_tab)
+
     # === EVENTS ===
-    gen1_btn.click(generate_caption_1, [image_input, system1, user1], out1)
-    gen2_btn.click(generate_caption_2, [image_input, system2, user2], out2)
-    gen3_btn.click(generate_caption_3, [image_input, system3, user3], out3)
+    gen1_btn.click(generate_caption, [image_input, system1, user1], out1)
+    gen2_btn.click(generate_caption, [image_input, system2, user2], out2)
+    gen3_btn.click(generate_caption, [image_input, system3, user3], out3)
     ask_btn.click(answer_question, [image_input, question_input], qa_output)
 
-    key_btn.click(lambda s1,u1,k: insert_key(s1,u1,k), [system1,user1,keywords_input], [system1,user1])
-    que_btn.click(lambda s1,u1,q: insert_que(s1,u1,q), [system1,user1,question_input], [system1,user1])
-    use_btn.click(lambda s1,u1,c: insert_use(s1,u1,c), [system1,user1,custom_instruction_input], [system1,user1])
-    not_btn.click(lambda s1,u1,a: insert_not(s1,u1,a), [system1,user1,avoid_input], [system1,user1])
+    # Template logic — update only current tab
+    def handle_template(btn_type, tab, s1, u1, s2, u2, s3, u3, k, c, q, a):
+        key_f, que_f, use_f, not_f = create_template_functions()
+        mapping = {
+            "key": key_f, "que": que_f, "use": use_f, "not": not_f
+        }
+        fn = mapping.get(btn_type)
+        if not fn:
+            return s1, u1, s2, u2, s3, u3
+        if tab == "casual":
+            s1, u1 = fn(s1, u1, k or c or q or a)
+        elif tab == "friendly":
+            s2, u2 = fn(s2, u2, k or c or q or a)
+        elif tab == "erotic":
+            s3, u3 = fn(s3, u3, k or c or q or a)
+        return s1, u1, s2, u2, s3, u3
+
+    for b, t in [(key_btn, "key"), (que_btn, "que"), (use_btn, "use"), (not_btn, "not")]:
+        b.click(
+            handle_template,
+            [gr.State(t), active_tab, system1, user1, system2, user2, system3, user3,
+             keywords_input, custom_instruction_input, question_input, avoid_input],
+            [system1, user1, system2, user2, system3, user3],
+        )
 
     def handle_export(k, c, a, q, c1, c2, c3, qa, img):
         msg, fd = export_joycaption_data(k, c, a, q, c1, c2, c3, qa, img)
         if fd:
             js, fn = fd
-            p = os.path.join(tempfile.gettempdir(), fn)
-            with open(p, "w", encoding="utf-8") as f: f.write(js)
-            return gr.update(value=msg, visible=True), gr.update(value=p, visible=True)
+            path = os.path.join(tempfile.gettempdir(), fn)
+            with open(path, "w", encoding="utf-8") as f:
+                f.write(js)
+            return gr.update(value=msg, visible=True), gr.update(value=path, visible=True)
         return gr.update(value=msg, visible=True), gr.update(visible=False)
 
-    export_btn.click(handle_export,
+    export_btn.click(
+        handle_export,
         [keywords_input, custom_instruction_input, avoid_input, question_input,
          out1, out2, out3, qa_output, image_input],
         [export_out, export_file]