Update app.py

app.py

@@ -1,723 +1,404 @@
-import spaces
-import gradio as gr
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import time
+"""
+Venice Edition NSFW Enhancer — consolidated single-call build.
+
+One @spaces.GPU invocation produces every output the blog needs:
+    1. Title            (click-bait, 8–14 words, SEO keywords woven in)
+    2. Meta description (140–160 chars, teaser one-liner)
+    3. H2 heading       (6–10 words, suggestive)
+    4. Rant             (120–200 word roasty / funny confession)
+    5. Tags             (each seed tag followed by 2 slang synonyms)
+
+Why one call? Each @spaces.GPU allocation has scheduling + warm-up overhead and
+consumes its own slice of the daily ZeroGPU quota. Three sequential buttons
+cost ≈60s of effective runtime per image; the consolidated function returns
+everything in ≈30–45s and uses ~half the quota.
+
+Prompts are intentionally short — the previous versions were multi-paragraph
+specs that blew up the input token budget for little benefit.
+"""
+
+import os
 import re
 import gc
-import os
-import shutil
 import json
-from datetime import datetime, timezone
 import uuid
-import random
+import shutil
+from datetime import datetime, timezone
+
+import spaces
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Storage optimization - redirect cache to temporary directories
-os.environ["HF_HOME"] = "/tmp/hf_cache"
+# ─── Cache redirection ──────────────────────────────────────────────────────
+os.environ["HF_HOME"]           = "/tmp/hf_cache"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/transformers_cache"
-os.environ["HF_DATASETS_CACHE"] = "/tmp/datasets_cache"
-os.environ["TORCH_HOME"] = "/tmp/torch_cache"
+os.environ["HF_DATASETS_CACHE"]  = "/tmp/datasets_cache"
+os.environ["TORCH_HOME"]         = "/tmp/torch_cache"
 
 MODEL_PATH = "dphn/dolphin-2.6-mistral-7b-dpo"
+print(f"🚀 Loading Venice Edition NSFW Enhancer v4.0 (consolidated)")
+print(f"📦 Model: {MODEL_PATH}")
 
 TITLE = """
-<div style="text-align: center; margin: 20px 0;">
-<h1>🌊 Venice Edition NSFW Enhancer</h1>
-<p><strong>Transform JoyCaption output into SEO-optimized, uncensored descriptions</strong></p>
-<p><em>Powered by Dolphin 2.6 Mistral 7B DPO - Venice Edition</em></p>
+<div style="text-align:center;margin:20px 0;">
+  <h1>🌊 Venice Edition NSFW Enhancer</h1>
+  <p><strong>One-shot: Title · Meta · H2 · Rant · Tags</strong></p>
+  <p><em>Dolphin 2.6 Mistral 7B DPO</em></p>
 </div>
 <hr>
 """
 
-print("🚀 Loading Venice Edition NSFW Enhancer... v3.5")
-print(f"📦 Model: {MODEL_PATH}")
+# ─── Housekeeping ───────────────────────────────────────────────────────────
 
-def cleanup_memory_aggressive():
+def _cleanup():
     try:
-        temp_dirs = ["/tmp/hf_cache", "/tmp/transformers_cache", "/tmp/datasets_cache", "/tmp/torch_cache"]
-        for temp_dir in temp_dirs:
-            if os.path.exists(temp_dir):
-                try:
-                    shutil.rmtree(temp_dir, ignore_errors=True)
-                except:
-                    pass
+        for d in ("/tmp/hf_cache", "/tmp/transformers_cache",
+                  "/tmp/datasets_cache", "/tmp/torch_cache"):
+            if os.path.exists(d):
+                shutil.rmtree(d, ignore_errors=True)
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
             torch.cuda.synchronize()
     except Exception as e:
-        print(f"⚠️ Cleanup warning: {e}")
+        print(f"⚠️ cleanup: {e}")
 
-def cleanup_storage():
-    cleanup_memory_aggressive()
-    print("✅ Storage and memory cleanup completed")
+# ─── JoyCaption JSON import ────────────────────────────────────────────────
 
-def process_joycaption_import(json_file):
-    """Process uploaded JoyCaption JSON (v6.0 sample schema) and return field data"""
+def import_joycaption_json(json_file):
+    """Parse a JoyCaption v6 export and pre-fill the captions / tags / URL fields."""
     if not json_file:
-        return "❌ No file uploaded", "", "", "", "", None
+        return "❌ No file uploaded", "", "", "", None
     try:
-        if hasattr(json_file, 'name'):
-            with open(json_file.name, 'r', encoding='utf-8') as f:
-                data = json.load(f)
-        else:
-            with open(json_file, 'r', encoding='utf-8') as f:
-                data = json.load(f)
-
-        if not isinstance(data, dict) or 'data' not in data:
-            return "❌ Invalid JSON format - missing 'data' field", "", "", "", "", None
-
-        joy_data = data['data']
-        success_msg = f"✅ Imported JoyCaption v6.0 at {data.get('timestamp', 'unknown time')}"
-
-        # Build captions (Friendly, Casual, Erotic all included)
-        desc = joy_data.get('descriptions', {}) or {}
-        captions_list = []
-        if desc.get('friendly'):
-            captions_list.append(f"═══ FRIENDLY ═══\n{desc['friendly']}")
-        if desc.get('casual'):
-            captions_list.append(f"═══ CASUAL ═══\n{desc['casual']}")
-        if desc.get('erotic'):
-            captions_list.append(f"═══ EROTIC ═══\n{desc['erotic']}")
-        all_captions = "\n\n".join(captions_list)
-        if captions_list:
-            all_captions = f"📊 IMPORTED {len(captions_list)} CAPTION TONES:\n\n" + all_captions
-
-        # Tags as keywords
-        keywords = (joy_data.get('tags') or "").strip()
-
-        # Corrections/context (user instructions)
+        path = json_file.name if hasattr(json_file, "name") else json_file
+        with open(path, "r", encoding="utf-8") as f:
+            raw = json.load(f)
+        data = raw.get("data") or {}
+        desc = data.get("descriptions") or {}
         parts = []
-        if joy_data.get('mention'):
-            parts.append(f"MENTION: {joy_data['mention']}")
-        if joy_data.get('avoid'):
-            parts.append(f"AVOID: {joy_data['avoid']}")
-        if joy_data.get('ask'):
-            parts.append(f"ASK: {joy_data['ask']}")
-        qa = joy_data.get('qa') or {}
-        if qa.get('question') or qa.get('answer'):
-            parts.append(f"Q&A: Q: {qa.get('question','')} A: {qa.get('answer','')}")
-        corrections = " | ".join(parts)
-
-        # Image URL
-        image_url = joy_data.get('image_path') or ""
-
-        import_details = []
-        if all_captions: import_details.append("Captions")
-        if keywords: import_details.append("Tags")
-        if corrections: import_details.append("Instructions/Q&A")
-        summary = f"{success_msg}\n📊 Imported: {', '.join(import_details) if import_details else 'No usable fields'}"
-
-        return summary, all_captions, keywords, corrections, image_url, data
-    except json.JSONDecodeError as e:
-        return f"❌ Invalid JSON file: {str(e)}", "", "", "", "", None
+        for tone in ("casual", "friendly", "erotic"):
+            if desc.get(tone):
+                parts.append(f"[{tone.upper()}]\n{desc[tone]}")
+        captions = "\n\n".join(parts)
+        tags     = (data.get("tags") or "").strip()
+        img_url  = (data.get("image_path") or "").strip()
+        return f"✅ Imported {len(parts)} tone(s)", captions, tags, img_url, raw
     except Exception as e:
-        return f"❌ Import failed: {str(e)}", "", "", "", "", None
+        return f"❌ Import failed: {e}", "", "", "", None
+
+# ─── Model load ────────────────────────────────────────────────────────────
+
+print("📦 Loading model + tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_PATH,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True,
+    low_cpu_mem_usage=True,
+)
+model.eval()
+print("✅ Model ready")
+_cleanup()
 
-print("📦 Loading Venice Edition model and tokenizer...")
-try:
-    tokenizer = AutoTokenizer.from_pretrained(
-        MODEL_PATH,
-        trust_remote_code=True
-    )
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_PATH,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-        trust_remote_code=True,
-        low_cpu_mem_usage=True
-    )
-    model.eval()
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-    print("✅ Venice Edition model loaded and ready!")
-    cleanup_storage()
-except Exception as e:
-    print(f"❌ Error loading model: {e}")
-    raise e
-
-def strip_markers(text: str) -> str:
-    if not text:
-        return ""
-    for marker in ['═══ FRIENDLY ═══','═══ CASUAL ═══','═══ EROTIC ═══','📊 IMPORTED','CAPTION TONES:']:
-        text = text.replace(marker, '')
-    return ' '.join(text.split())
+# ─── Prompt builders (short and opinionated) ───────────────────────────────
 
-def build_keywords_hint(seo_keywords: str) -> str:
-    if not seo_keywords or not seo_keywords.strip():
-        return ""
-    ks = [k.strip() for k in re.split(r"[,\n]+", seo_keywords) if k.strip()]
-    return ", ".join(ks)
-
-def _truncate_words(text: str, max_words: int) -> str:
-    words = re.findall(r"\S+", text or "")
-    if len(words) <= max_words:
-        return text
-    return " ".join(words[:max_words]).rstrip(" ,;.")
-
-# Default editable prompts for BLOG POST (HTML, longer, late-night confession style)
-DEFAULT_SYSTEM_PROMPT = (
-    "You are an adult-themed blog writer. You write sensual, suggestive, late-night confession style posts "
-    "focused on mood, atmosphere, inner thoughts, and fantasies. "
-    "Avoid graphic descriptions of sexual acts or explicit anatomy. "
-    "Write like a real person: natural flow, varied sentence lengths, sensory detail, contractions, vivid imagery. "
-    "NEVER use the words 'woman' or 'women'—instead use slang like babe, chick, beauty, stunner, hottie, goddess, vixen, etc. "
-    "No meta-commentary. Write as a blog post, not an image caption."
+SYSTEM_PROMPT = (
+    "You write adult-blog content. Voice: suggestive, cheeky, roasty, funny. "
+    "Never use 'woman' or 'women' — swap in slang (babe, chick, stunner, "
+    "hottie, vixen, goddess, bombshell). No graphic anatomy. "
+    "Return ONLY the 5 labelled sections below, nothing else, no preamble."
 )
 
-DEFAULT_USER_PROMPT = """You are given detailed JoyCaption descriptions and some priority rules.
+USER_TEMPLATE = """IMAGE DESCRIPTIONS:
+{captions}
+
+TAG SEEDS:
+{tags}
 
-JOY CAPTION DESCRIPTIONS (SOURCE MATERIAL):
-[[DESCRIPTIONS]]
+Write each section. Follow the labels exactly.
 
-PRIORITY RULES (follow these if present):
-[[RULES]]
+===TITLE===
+<click-bait blog title, 8–14 words, weave in 1–2 tag keywords>
 
-Write a first-person, late-night confession style blog post (about 220–350 words) as HTML.
+===META===
+<SEO meta description, 1–2 sentences, 140–160 characters total, teasing>
 
-Tone and content:
-- Adult, sensual, suggestive
-- Not a dry description of the image; instead, inner thoughts, fantasies, micro-stories, atmosphere
-- Focus on feelings, sensations, impressions, and imagined scenarios inspired by the descriptions
-- Avoid graphic descriptions of sexual acts or explicit anatomy
-- NEVER use "woman" or "women"—use slang: babe, chick, beauty, stunner, hottie, goddess, vixen, etc.
+===H2===
+<single H2 heading, 6–10 words, suggestive, no quotes>
 
-Structure and formatting:
-- Output VALID HTML only
-- Start with a single <h2> heading that naturally includes 1–2 of the SEO KEYWORDS from the rules
-- Then write 2–4 paragraphs of confession-style story text
-- Whenever you naturally use any of the SEO KEYWORDS from the rules (the line starting with "KEYWORDS"), wrap them in <strong>...</strong> tags
-- Do NOT include <html>, <head>, <body>, or any other boilerplate; only <h2> and paragraphs
+===RANT===
+<roasty, funny, slightly savage 120–200 word rant riffing on the scene.
+First-person narrator. Adult tone but not graphic. Wrap any tag keyword
+you use in <strong>…</strong>.>
 
-Return ONLY the HTML blog post (the <h2> heading + paragraphs), no explanation or extra text:
+===TAGS===
+<for each seed tag, output the tag followed by 2 slang synonyms.
+Comma-separated, no hashtags, no duplicates. Example:
+beach, shoreline, coastal spot, sunset, golden hour, dusk glow>
 """
 
-def build_priority_rules(instructions: str, keywords: str) -> str:
-    """Build priority rules string from instructions (MENTION/AVOID/ASK/Q&A) and keywords."""
-    rules = []
-    if instructions and instructions.strip():
-        # Parse MENTION, AVOID, ASK, Q&A
-        mention_match = re.search(r'MENTION:\s*([^|]+)', instructions)
-        if mention_match:
-            rules.append(f"MENTION: {mention_match.group(1).strip()}")
-        
-        keywords_str = build_keywords_hint(keywords)
-        if keywords_str:
-            rules.append(f"KEYWORDS (integrate naturally): {keywords_str}")
-        
-        qa_match = re.search(r'Q&A:\s*Q:\s*([^A]+)A:\s*([^|]+)', instructions)
-        if qa_match:
-            rules.append(f"QUESTION/ANSWER: {qa_match.group(1).strip()} / {qa_match.group(2).strip()}")
-        
-        avoid_match = re.search(r'AVOID:\s*([^|]+)', instructions)
-        if avoid_match:
-            rules.append(f"AVOID: {avoid_match.group(1).strip()}")
-    elif keywords and keywords.strip():
-        # If no instructions but keywords present
-        keywords_str = build_keywords_hint(keywords)
-        rules.append(f"KEYWORDS (integrate naturally): {keywords_str}")
-    
-    return "\n".join(rules) if rules else "(none)"
-
-def fill_user_prompt_template(user_prompt_tpl: str, descriptions: str, rules: str) -> str:
-    tpl = user_prompt_tpl or ""
-    tpl = tpl.replace("[[DESCRIPTIONS]]", descriptions or "(none)")
-    tpl = tpl.replace("[[RULES]]", rules or "(none)")
-    return tpl
-
-@spaces.GPU(duration=65)
-@torch.no_grad()
-def enhance_text(all_captions, seo_keywords="", user_corrections="", sys_override="", user_override=""):
-    """
-    Generate an adult, sensual HTML blog post (≈220–350 words) using
-    full JoyCaption descriptions + priority rules, with <h2> heading and <strong>KEYWORDS</strong>.
-    """
-    src = strip_markers(all_captions or "")
-    if not src:
-        return "❌ Please provide JoyCaption descriptions first (via JSON import or paste)."
-    try:
-        system_prompt = (sys_override or DEFAULT_SYSTEM_PROMPT)
-        priority_rules = build_priority_rules(user_corrections, seo_keywords)
-        user_prompt = fill_user_prompt_template(
-            user_override or DEFAULT_USER_PROMPT,
-            src, priority_rules
-        )
-        full = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
-        inputs = tokenizer(full, return_tensors="pt", truncation=True, max_length=3000, padding=True)
-        device = next(model.parameters()).device
-        inputs = {k: v.to(device) for k, v in inputs.items()}
-        # Higher randomness and length for more varied, longer posts
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=700,
-            temperature=0.95,
-            top_p=0.96,
-            top_k=80,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id,
-            repetition_penalty=1.05,
-            use_cache=False
-        )
-        ilen = inputs["input_ids"].shape[1]
-        gen = outputs[0][ilen:]
-        txt = tokenizer.decode(gen, skip_special_tokens=True, clean_up_tokenization_spaces=True)
-
-        # Clean special tokens and boilerplate
-        txt = re.sub(r'<\|im_.*?\|>', '', txt, flags=re.DOTALL).strip()
-        txt = re.sub(r'^(assistant|Assistant)\s*', '', txt).strip()
-        txt = re.sub(r'^\s*:', '', txt).strip()
-
-        # Remove any full-page boilerplate, keep from first <h2> onward
-        lower = txt.lower()
-        idx = lower.find("<h2")
-        if idx != -1:
-            txt = txt[idx:]
-        # Strip trailing </body>, </html>, etc.
-        txt = re.sub(r'(?is)</body>.*$', '', txt).strip()
-        txt = re.sub(r'(?is)</html>.*$', '', txt).strip()
-
-        # If still no <h2>, wrap the first line in <h2> as a fallback
-        if "<h2" not in txt.lower():
-            lines = txt.splitlines()
-            if lines:
-                first = lines[0].strip()
-                rest = "\n".join(lines[1:]).strip()
-                if first:
-                    txt = f"<h2>{first}</h2>\n\n{rest}"
-
-        txt = re.sub(r"\s{2,}", " ", txt)
-        cleanup_memory_aggressive()
-        cleanup_storage()
-        return txt if txt else "❌ No blog post generated"
-    except Exception as e:
-        cleanup_memory_aggressive()
-        cleanup_storage()
-        return f"❌ Error: {str(e)[:200]}..."
-
-def create_title_prompt(descriptions: str, instructions: str, seo_keywords: str):
-    kw_hint = build_keywords_hint(seo_keywords)
-    system_prompt = (
-        "You write short, catchy, SEO-optimized, click-bait titles for adult-themed blog posts. "
-        "Titles should be around 8–14 words, very enticing, and include 1–2 provided keywords naturally. "
-        "Make them bold, suggestive, and optimized for adult search traffic. "
-        "Avoid graphic descriptions of sexual acts or explicit anatomy. "
-        "NEVER use 'woman' or 'women'—use slang like babe, chick, beauty, stunner, hottie, etc. "
-        "Return ONLY the title text (no quotes/numbering)."
+def _build_prompt(captions: str, tags: str) -> str:
+    cap = (captions or "").strip() or "(none provided)"
+    tag = (tags or "").strip() or "(none)"
+    user = USER_TEMPLATE.format(captions=cap, tags=tag)
+    return (
+        f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
+        f"<|im_start|>user\n{user}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
     )
-    user_prompt = f"""Write ONE click-bait adult blog title (8–14 words) for the post.
-
-JOY CAPTION DESCRIPTIONS (SOURCE MATERIAL):
-{descriptions or '(none provided)'}
-
-INSTRUCTIONS (if present):
-{instructions or '(none)'}
-
-KEYWORDS (include 1–2 naturally for SEO):
-{kw_hint or '(none)'}
-
-Rules:
-- 8–14 words
-- Sentence-style phrase (not all caps, no numbering)
-- No quotes, no leading numbers
-- NEVER use "woman" or "women"—use slang: babe, chick, beauty, stunner, hottie, goddess, vixen, etc.
-- Suggestive and enticing, but not graphically explicit
-- Include 1–2 keywords naturally
-- SEO-optimized and highly click-bait
-Return only the title text:"""
-    return system_prompt, user_prompt
-
-def _normalize_title_sentence(s: str, keywords: str) -> str:
-    if not s:
-        return s
-    s = re.sub(r"[.,;:!?\"'(){}\[\]<>/\\\-–—•…]+", " ", s)
-    s = re.sub(r"\s{2,}", " ", s).strip()
-    s = s.capitalize()
-    words = s.split()
-    # Allow longer titles before truncation
-    if len(words) > 14:
-        s = " ".join(words[:14])
-    if keywords and keywords.strip():
-        ks = [k.strip().lower() for k in re.split(r"[,\n]+", keywords) if k.strip()]
-        if ks and not any(k in s.lower() for k in ks):
-            if len(words) < 14:
-                s = f"{s} {ks[0]}"
-    return s
-
-@spaces.GPU(duration=30)
-@torch.no_grad()
-def generate_single_title(all_captions, user_corrections="", seo_keywords=""):
-    """
-    Generate one short, suggestive, click-bait blog title (8–14 words);
-    SEO-optimized with keywords; natural sentence, no quotes/numbering.
-    """
-    src = strip_markers(all_captions or "")
-    if not src.strip():
-        return "❌ Please provide JoyCaption descriptions first (via JSON import or paste)."
-    try:
-        system_prompt, user_prompt = create_title_prompt(src, user_corrections, seo_keywords)
-        full = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
-        inputs = tokenizer(full, return_tensors="pt", truncation=True, max_length=2500, padding=True)
-        device = next(model.parameters()).device
-        inputs = {k: v.to(device) for k, v in inputs.items()}
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=120,
-            temperature=0.95,
-            top_p=0.95,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id,
-            use_cache=False
-        )
-        ilen = inputs["input_ids"].shape[1]
-        gen = outputs[0][ilen:]
-        txt = tokenizer.decode(gen, skip_special_tokens=True, clean_up_tokenization_spaces=True)
-        txt = re.sub(r'<\|im_.*?\|>', '', txt).strip()
-        txt = re.sub(r'^(assistant|Assistant)\s*', '', txt).strip()
-        if "\n" in txt:
-            for line in txt.splitlines():
-                if line.strip():
-                    txt = line.strip()
-                    break
-        txt = _normalize_title_sentence(txt, seo_keywords)
-        cleanup_memory_aggressive()
-        cleanup_storage()
-        return txt if txt else "❌ No title generated"
-    except Exception as e:
-        cleanup_memory_aggressive()
-        cleanup_storage()
-        return f"❌ Error: {str(e)[:200]}..."
 
-BANNED_ACTIONS = {
-    "kissing","hugging","holding","touching","walking","running","posing","smiling","standing","looking",
-    "talking","laughing","dancing","drinking","eating","singing","sleeping","licking","sucking","riding"
-}
+# ─── Output parser ─────────────────────────────────────────────────────────
+
+_SECTION_RE = re.compile(
+    r"===\s*(TITLE|META|H2|RANT|TAGS)\s*===\s*(.*?)(?=(?:===\s*(?:TITLE|META|H2|RANT|TAGS)\s*===)|\Z)",
+    re.DOTALL | re.IGNORECASE,
+)
 
-@spaces.GPU(duration=25)
+def _parse_sections(raw: str) -> dict:
+    out = {"title": "", "meta": "", "h2": "", "rant": "", "tags": ""}
+    if not raw:
+        return out
+    for m in _SECTION_RE.finditer(raw):
+        key = m.group(1).lower()
+        body = m.group(2).strip()
+        # Strip the placeholder brackets the model sometimes echoes
+        body = re.sub(r"^<|>$", "", body).strip()
+        if key == "title" or key == "h2":
+            body = body.splitlines()[0].strip(' "\'') if body else ""
+        out[key if key != "h2" else "h2"] = body
+    return out
+
+def _postprocess(parts: dict, tags_seed: str) -> dict:
+    # Title: keep to 14 words, no quotes, capitalize
+    t = parts.get("title", "")
+    t = re.sub(r'["\']', "", t).strip()
+    words = t.split()
+    if len(words) > 14:
+        t = " ".join(words[:14])
+    parts["title"] = t[:140]
+
+    # Meta: squash whitespace, clip to 180 chars
+    m = re.sub(r"\s+", " ", parts.get("meta", "")).strip()
+    parts["meta"] = m[:180]
+
+    # H2: one line, no tags
+    h2 = re.sub(r"<[^>]+>", "", parts.get("h2", "")).strip()
+    parts["h2"] = h2
+
+    # Rant: strip any leading/trailing code fences, collapse blank runs
+    r = parts.get("rant", "")
+    r = re.sub(r"^```.*?\n|```\s*$", "", r, flags=re.DOTALL)
+    r = re.sub(r"\n{3,}", "\n\n", r).strip()
+    parts["rant"] = r
+
+    # Tags: backfill seeds if the model skipped them
+    seeds = [s.strip() for s in re.split(r"[,\n]+", tags_seed or "") if s.strip()]
+    tagstr = parts.get("tags", "")
+    tagstr = re.sub(r"#", "", tagstr)
+    tagstr = re.sub(r"\s+", " ", tagstr)
+    tag_list = [x.strip() for x in tagstr.split(",") if x.strip()]
+    seen, ordered = set(), []
+    for seed in seeds:
+        if seed.lower() not in seen:
+            seen.add(seed.lower()); ordered.append(seed)
+    for t in tag_list:
+        if t.lower() not in seen:
+            seen.add(t.lower()); ordered.append(t)
+    parts["tags"] = ", ".join(ordered)
+    return parts
+
+# ─── The one GPU call ──────────────────────────────────────────────────────
+
+@spaces.GPU(duration=55)
 @torch.no_grad()
-def generate_keywords_expanded(tags_text, all_captions="", user_corrections=""):
+def generate_all(captions: str, tags: str):
     """
-    For each tag, ask the model for adult-blog synonyms (lexical, precise),
-    then keep 2–3 short expressions per tag.
+    Single forward pass → all five fields.
+    Returns (title, meta, h2, rant, tags) so Gradio can fan them out.
     """
-    if not tags_text or not tags_text.strip():
-        return "❌ No tags found in JSON."
-    tags = [t.strip() for t in re.split(r"[,\n]+", tags_text) if t.strip()]
-    if not tags:
-        return "❌ No valid tags."
+    if not (captions and captions.strip()):
+        empty = "❌ Provide JoyCaption descriptions first."
+        return empty, "", "", "", ""
     try:
-        system_prompt = (
-            "You are an SEO assistant for an adult-themed blog. For EACH tag provided, produce 3–5 close lexical synonyms/phrases "
-            "(prefer nouns/adjectives, 1–2 words each). Do NOT invent actions, relationships, or scenes unless the tag itself is an action. "
-            "Stay precise, adult-context-appropriate, and avoid hashtags/quotes/digits/duplicates. Output one line per tag."
-        )
-        tag_lines = "\n".join([f"- {t}" for t in tags])
-        instr = f"\nCONTEXT (PRIORITY): {user_corrections}" if user_corrections and user_corrections.strip() else ""
-        cap_hint = f"\nTONES HINT: {strip_markers(all_captions)[:250]}" if all_captions and all_captions.strip() else ""
-        user_prompt = f"""Tags:
-{tag_lines}
-For each tag, return a single line strictly in the format (3–5 options):
-tag: expr1, expr2, expr3[, expr4][, expr5]
-Rules:
-- Adult blog tone
-- Prefer nouns/adjectives; 1–2 words each
-- No hashtags, no digits, no quotes
-- Do NOT invent actions unless the tag itself is an action
-{instr}{cap_hint}
-Only return the list, no extra text."""
-        full = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
-        inputs = tokenizer(full, return_tensors="pt", truncation=True, max_length=1700, padding=True)
+        prompt = _build_prompt(captions, tags)
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True,
+                           max_length=2600, padding=True)
         device = next(model.parameters()).device
         inputs = {k: v.to(device) for k, v in inputs.items()}
-        outputs = model.generate(
+        output = model.generate(
             **inputs,
-            max_new_tokens=260,
-            temperature=0.6,
-            top_p=0.92,
+            max_new_tokens=900,
+            temperature=0.9,
+            top_p=0.95,
             do_sample=True,
+            repetition_penalty=1.05,
             pad_token_id=tokenizer.eos_token_id,
             eos_token_id=tokenizer.eos_token_id,
-            use_cache=False
+            use_cache=True,
         )
-        ilen = inputs["input_ids"].shape[1]
-        gen = outputs[0][ilen:]
-        raw = tokenizer.decode(gen, skip_special_tokens=True, clean_up_tokenization_spaces=True)
-        raw = re.sub(r'<\|im_.*?\|>', '', raw).strip()
-
-        flat = []
-        found = {}
-        for line in raw.splitlines():
-            line = line.strip()
-            if not line or ":" not in line:
-                continue
-            tag, syns = line.split(":", 1)
-            tag = tag.strip().strip("-•").strip()
-            if not tag:
-                continue
-            syns_list = [s.strip().strip(' "\'') for s in syns.split(",") if s.strip()]
-            cleaned = []
-            for s in syns_list:
-                if not s or any(ch.isdigit() for ch in s) or "#" in s:
-                    continue
-                words = s.split()
-                if not (1 <= len(words) <= 2):
-                    continue
-                if s.lower() == tag.lower():
-                    continue
-                if any(w.lower() in BANNED_ACTIONS for w in words):
-                    continue
-                if any(w.lower().endswith("ing") for w in words) and tag.lower() not in BANNED_ACTIONS:
-                    continue
-                cleaned.append(s)
-            seen = set()
-            dedup = []
-            for s in cleaned:
-                key = s.lower()
-                if key not in seen:
-                    seen.add(key)
-                    dedup.append(s)
-            k = 3 if random.random() < 0.5 else 2
-            if len(dedup) >= k:
-                chosen = random.sample(dedup, k)
-            elif len(dedup) >= 2:
-                chosen = random.sample(dedup, 2)
-            else:
-                chosen = dedup
-            found[tag.lower()] = [tag] + chosen
-
-        for t in tags:
-            key = t.lower()
-            if key not in found:
-                found[key] = [t]
-
-        for t in tags:
-            flat.extend(found.get(t.lower(), [t]))
-
-        del inputs, outputs
-        cleanup_memory_aggressive()
-        cleanup_storage()
-        out_list = []
-        seen_all = set()
-        for x in flat:
-            key = x.lower()
-            if key not in seen_all:
-                seen_all.add(key)
-                out_list.append(x)
-        return ", ".join(out_list)
+        gen = output[0][inputs["input_ids"].shape[1]:]
+        raw = tokenizer.decode(gen, skip_special_tokens=True).strip()
+        raw = re.sub(r"<\|im_.*?\|>", "", raw).strip()
+        parts = _postprocess(_parse_sections(raw), tags)
+        _cleanup()
+        return parts["title"], parts["meta"], parts["h2"], parts["rant"], parts["tags"]
     except Exception as e:
-        cleanup_memory_aggressive()
-        cleanup_storage()
-        return f"❌ Error: {str(e)[:200]}..."
+        _cleanup()
+        err = f"❌ Error: {str(e)[:200]}"
+        return err, "", "", "", ""
 
-def _first_title_line(titles_text: str) -> str:
-    if not titles_text:
-        return ""
-    s = titles_text.splitlines()[0].strip()
-    s = re.sub(r'^\s*\d+\.\s*', '', s).strip(' "\'')
-    return s
+# ─── Export ────────────────────────────────────────────────────────────────
 
-def export_blog_json(raw_data, title_text, caption_text, keywords_text, image_url):
-    """
-    Export the blog data as a single JSON file for download.
-    """
-    if not raw_data:
+def export_blog_json(raw_data, title, meta, h2, rant, tags, image_url):
+    if not (title or rant):
         return None
-    out = {
+    payload = {
         "timestamp": datetime.now(timezone.utc).isoformat(),
-        "source": "Venice Edition Blog Generator",
-        "input": {
-            "tags": (raw_data.get("data", {}) or {}).get("tags"),
-            "mention": (raw_data.get("data", {}) or {}).get("mention"),
-            "avoid": (raw_data.get("data", {}) or {}).get("avoid"),
-            "ask": (raw_data.get("data", {}) or {}).get("ask"),
-            "qa": (raw_data.get("data", {}) or {}).get("qa"),
-            "image_path": (raw_data.get("data", {}) or {}).get("image_path"),
-        },
+        "source":    "Venice Edition v4.0",
         "output": {
-            "image_url": (image_url or (raw_data.get("data", {}) or {}).get("image_path") or "").strip(),
-            "title": _first_title_line(title_text or ""),
-            "keywords": [k.strip() for k in (keywords_text or "").split(",") if k.strip()],
-            "post": (caption_text or "").strip(),
-        }
+            "image_url":   (image_url or "").strip(),
+            "title":       title or "",
+            "meta":        meta or "",
+            "h2_heading":  h2 or "",
+            "rant_html":   rant or "",
+            "tags":        [t.strip() for t in (tags or "").split(",") if t.strip()],
+        },
+        "source_import": raw_data,
     }
     path = f"/tmp/blog_{uuid.uuid4().hex}.json"
     with open(path, "w", encoding="utf-8") as f:
-        json.dump(out, f, ensure_ascii=False, indent=2)
+        json.dump(payload, f, ensure_ascii=False, indent=2)
     return path
 
-# UI
+# ─── Gradio UI ─────────────────────────────────────────────────────────────
+
 with gr.Blocks(title="Venice Edition NSFW Enhancer", theme=gr.themes.Soft()) as demo:
     gr.HTML(TITLE)
-
-    raw_import_state = gr.State(value=None)
+    raw_state = gr.State(value=None)
 
     with gr.Row():
-        # LEFT COLUMN
+        # ── Left: inputs ────────────────────────────────────────────────
         with gr.Column(scale=1):
-            gr.Markdown("### 📥 Import JoyCaption Data")
+            gr.Markdown("### 📥 Import JoyCaption JSON (optional)")
             with gr.Row():
-                import_file = gr.File(label="Upload JoyCaption JSON", file_types=[".json"], scale=2)
-                import_btn = gr.Button("📥 Import", variant="primary", scale=1, size="lg")
+                import_file = gr.File(label="JoyCaption JSON", file_types=[".json"], scale=2)
+                import_btn  = gr.Button("📥 Import", variant="primary", scale=1, size="lg")
+            import_status = gr.Textbox(label="Status", lines=2, interactive=False, visible=False)
 
-            import_status = gr.Textbox(label="Import Status", lines=3, interactive=False, visible=False)
-
-            gr.Markdown("### 📝 Caption Data (All Tones)")
+            gr.Markdown("### 📝 Inputs")
             all_captions_input = gr.Textbox(
-                placeholder="Import from JoyCaption or paste all caption tones here...",
                 label="JoyCaption Descriptions",
-                lines=12,
-                max_lines=18,
-                info="Shows caption tones from JoyCaption (Friendly, Casual, Erotic) — may be 1–3 tones"
-            )
-
-            user_corrections_input = gr.Textbox(
-                placeholder="Your instructions from JSON (MENTION/AVOID/ASK/Q&A) or custom edits",
-                label="✏️ Your Instructions (PRIORITY)",
-                lines=4,
-                info="Your instructions override AI descriptions"
+                lines=12, max_lines=18,
+                placeholder="Paste JoyCaption output (one or more tones)...",
+                elem_id="venice_captions_input",
             )
-
             seo_keywords_input = gr.Textbox(
-                placeholder="Will auto-fill from JSON 'tags'",
-                label="🏷️ SEO Keywords (from JSON tags)",
+                label="🏷️ Tag Seeds",
                 lines=2,
-                info="We will expand each tag with 2–3 precise adult blog synonyms"
+                placeholder="comma,separated,seed,tags",
+                elem_id="venice_tags_input",
+            )
+            image_url_input = gr.Textbox(
+                label="🖼️ Image URL",
+                lines=1,
+                placeholder="https://...",
+                elem_id="venice_image_url",
             )
+            image_preview = gr.HTML(elem_id="venice_image_preview")
 
-            # Keywords section
-            keywords_output = gr.Textbox(
-                label="🏷️ Blog Keywords (Original + 2–3 Synonyms Each)",
-                lines=4,
-                max_lines=10,
-                show_copy_button=True,
-                placeholder="Expanded keyword list will appear here...",
-                info="Comma-separated list for an adult blog"
+            generate_btn = gr.Button(
+                "🌊 Generate All", variant="primary", size="lg",
+                elem_id="venice_generate_btn",
             )
-            generate_keywords_btn = gr.Button("🔑 Generate Keywords + Synonyms", variant="primary", size="lg")
 
-        # RIGHT COLUMN
+        # ── Right: outputs ──────────────────────────────────────────────
         with gr.Column(scale=1):
-            image_url_input = gr.Textbox(
-                placeholder="Paste image URL here...",
-                label="🖼️ Image URL",
-                lines=1
+            title_output = gr.Textbox(
+                label="📰 Title",
+                lines=2, show_copy_button=True,
+                elem_id="venice_title_output",
             )
-            gr.Markdown("📸 Blog Image Preview & Download")
-            image_preview = gr.HTML()
-
-            # Editable prompts (Blog Post)
-            gr.Markdown("### 🪄 Editable Prompts (Blog Post)")
-            system_prompt_input = gr.Textbox(
-                label="System Prompt (Blog Post)",
-                value=DEFAULT_SYSTEM_PROMPT,
-                lines=5
+            meta_output = gr.Textbox(
+                label="🔎 SEO Meta Description (≤160 chars)",
+                lines=2, show_copy_button=True,
+                elem_id="venice_meta_output",
             )
-            user_prompt_input = gr.Textbox(
-                label="User Prompt (Blog Post) — use [[DESCRIPTIONS]] [[RULES]]",
-                value=DEFAULT_USER_PROMPT,
-                lines=12
+            h2_output = gr.Textbox(
+                label="🏷️ H2 Heading",
+                lines=2, show_copy_button=True,
+                elem_id="venice_h2_output",
             )
-
-            title_output = gr.Textbox(
-                label="📰 Blog Post Title (8–14 words, SEO-optimized, click-bait, no 'woman/women')",
-                lines=2,
-                max_lines=4,
-                show_copy_button=True,
-                placeholder="Short suggestive blog title will appear here..."
+            rant_output = gr.Textbox(
+                label="🔥 Roasty Rant (120–200 words)",
+                lines=12, max_lines=20, show_copy_button=True,
+                elem_id="venice_rant_output",
             )
-            generate_title_btn = gr.Button("🎯 Generate Title", variant="primary", size="lg")
-
-            description_output = gr.Textbox(
-                label="🌊 HTML Blog Post (220–350 words, <h2> + <strong>keywords</strong>, no 'woman/women')",
-                lines=14,
-                max_lines=22,
-                show_copy_button=True,
-                placeholder="Your HTML blog post will appear here...",
-                info="Adult, sensual, confession-style HTML blog post; starts with <h2> heading using 1–2 SEO keywords, then paragraphs. Any used SEO keywords are wrapped in <strong>...</strong>."
+            tags_output = gr.Textbox(
+                label="🏷️ SEO Tags (seeds + slang synonyms)",
+                lines=4, show_copy_button=True,
+                elem_id="venice_tags_output",
             )
-            enhance_btn = gr.Button("🌊 Generate Blog Post", variant="primary", size="lg")
 
-            # Export blog JSON
-            export_file = gr.File(
-                label="Download Blog JSON",
-                interactive=False
-            )
-            export_btn = gr.Button("💾 Export Blog JSON", variant="secondary", size="lg")
-
-    # Handlers
-    def handle_import(json_file):
-        status, all_captions, keywords, corrections, image_url, raw = process_joycaption_import(json_file)
-        if image_url:
-            u = image_url.strip()
-            img_html = f"""
-<div>
-  <img src="{u}" style="max-width:100%; height:auto; border-radius:8px" />
-  <div style="margin-top:8px;">
-    <a href="{u}" download="image.jpg" target="_blank">⬇️ Download image (image.jpg)</a>
-  </div>
-</div>
-"""
+            export_btn  = gr.Button("💾 Export Blog JSON", variant="secondary")
+            export_file = gr.File(label="Download", interactive=False)
+
+    # ── Handlers ─────────────────────────��──────────────────────────────
+
+    def _handle_import(file):
+        status, captions, tags, url, raw = import_joycaption_json(file)
+        if url:
+            img = f'<img src="{url}" style="max-width:100%;border-radius:8px" />'
         else:
-            img_html = ""
+            img = ""
         return (
             gr.update(value=status, visible=True),
-            gr.update(value=all_captions),
-            gr.update(value=corrections),
-            gr.update(value=keywords),
-            gr.update(value=image_url),
-            gr.update(value=img_html),
-            raw
+            gr.update(value=captions),
+            gr.update(value=tags),
+            gr.update(value=url),
+            gr.update(value=img),
+            raw,
         )
 
     import_btn.click(
-        handle_import,
+        _handle_import,
         inputs=[import_file],
-        outputs=[import_status, all_captions_input, user_corrections_input, seo_keywords_input, image_url_input, image_preview, raw_import_state]
+        outputs=[import_status, all_captions_input, seo_keywords_input,
+                 image_url_input, image_preview, raw_state],
     )
 
-    def update_image_preview(url):
-        if url and url.strip() and (url.startswith("http://") or url.startswith("https://")):
-            u = url.strip()
-            return f"""
-<div>
-  <img src="{u}" style="max-width:100%; height:auto; border-radius:8px" />
-  <div style="margin-top:8px;">
-    <a href="{u}" download="image.jpg" target="_blank">⬇️ Download image (image.jpg)</a>
-  </div>
-</div>
-"""
+    def _update_preview(url):
+        u = (url or "").strip()
+        if u.startswith(("http://", "https://")):
+            return f'<img src="{u}" style="max-width:100%;border-radius:8px" />'
         return ""
 
-    image_url_input.change(update_image_preview, inputs=image_url_input, outputs=image_preview)
-
-    generate_title_btn.click(
-        generate_single_title,
-        inputs=[all_captions_input, user_corrections_input, seo_keywords_input],
-        outputs=title_output,
-        show_progress=True
-    )
-
-    enhance_btn.click(
-        enhance_text,
-        inputs=[all_captions_input, seo_keywords_input, user_corrections_input, system_prompt_input, user_prompt_input],
-        outputs=description_output,
-        show_progress=True
-    )
-
-    generate_keywords_btn.click(
-        generate_keywords_expanded,
-        inputs=[seo_keywords_input, all_captions_input, user_corrections_input],
-        outputs=keywords_output,
-        show_progress=True
+    image_url_input.change(_update_preview, inputs=image_url_input, outputs=image_preview)
+
+    # Clear old outputs on every fresh run → kills stale "Error" badges.
+    def _clear_outputs():
+        return "", "", "", "", ""
+
+    generate_btn.click(
+        _clear_outputs, inputs=None,
+        outputs=[title_output, meta_output, h2_output, rant_output, tags_output],
+        queue=False,
+    ).then(
+        generate_all,
+        inputs=[all_captions_input, seo_keywords_input],
+        outputs=[title_output, meta_output, h2_output, rant_output, tags_output],
+        show_progress=True,
     )
 
     export_btn.click(
         export_blog_json,
-        inputs=[raw_import_state, title_output, description_output, keywords_output, image_url_input],
+        inputs=[raw_state, title_output, meta_output, h2_output, rant_output,
+                tags_output, image_url_input],
         outputs=export_file,
-        show_progress=True
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()