Update app.py

app.py

@@ -2,18 +2,29 @@ import gradio as gr
 import torch
 import spaces
 from diffusers import StableDiffusionXLPipeline, DPMSolverMultistepScheduler
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, InferenceClient
 import random
-import re
+import os
 
-# ── Model — CyberIllustrious V8 ───────────────────────────────────────────────
+# ── HF Inference client (prompt expansion LLM) ────────────────────────────────
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+llm_client = InferenceClient(
+    model="mistralai/Mistral-7B-Instruct-v0.3",
+    token=HF_TOKEN,
+)
+
+# ── Image model — CyberIllustrious ────────────────────────────────────────────
 MODEL_REPO = "cyberdelia/latest_sdxl_models"
 MODEL_FILE = "CyberIllustrious_V8.0alt.safetensors"
 IL_POS     = "masterpiece, best quality, very aesthetic, absurdres, "
 IL_NEG     = "worst quality, low quality, bad quality, ugly, "
 
-print("Downloading CyberIllustrious V8...")
-local_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
+print("Downloading CyberIllustrious...")
+local_path = hf_hub_download(
+    repo_id=MODEL_REPO,
+    filename=MODEL_FILE,
+    token=HF_TOKEN,
+)
 print("Loading pipeline...")
 pipe = StableDiffusionXLPipeline.from_single_file(local_path, torch_dtype=torch.float16)
 pipe.scheduler = DPMSolverMultistepScheduler.from_config(
@@ -22,145 +33,76 @@ pipe.scheduler = DPMSolverMultistepScheduler.from_config(
 pipe.enable_attention_slicing()
 print("Ready.")
 
-# ── Smart prompt expansion (Fooocus-style) ─────────────────────────────────────
-# Detects scene type and expands with the right vocabulary.
-# Specific unique details (numbers, colours, states) get attention weights.
-
-SCENE_TAGS = {
-    "portrait": {
-        "keywords": ["woman", "man", "girl", "boy", "person", "face", "portrait",
-                     "lady", "guy", "model", "character"],
-        "pos": "sharp focus on face, detailed eyes, professional photography, "
-               "85mm lens, f/1.8 bokeh, skin texture, natural lighting, ",
-        "neg": "wide angle distortion, bad eyes, asymmetric face, blurry face, ",
-    },
-    "architecture": {
-        "keywords": ["building", "house", "tower", "city", "street", "bridge",
-                     "castle", "cathedral", "skyscraper", "facade", "structure"],
-        "pos": "architectural photography, sharp geometry, detailed textures, "
-               "golden hour lighting, wide angle lens, high resolution, ",
-        "neg": "distorted perspective, blurry, watermarks, ",
-    },
-    "landscape": {
-        "keywords": ["mountain", "forest", "ocean", "beach", "valley", "sky",
-                     "field", "river", "lake", "landscape", "nature", "countryside"],
-        "pos": "landscape photography, epic vista, golden hour, volumetric light, "
-               "sharp foreground, atmospheric perspective, 16mm lens, ",
-        "neg": "oversaturated, blurry horizon, flat lighting, ",
-    },
-    "interior": {
-        "keywords": ["room", "interior", "bedroom", "kitchen", "office",
-                     "living room", "hallway", "bathroom", "studio", "cafe"],
-        "pos": "interior photography, ambient lighting, detailed surfaces, "
-               "realistic materials, depth of field, architectural digest style, ",
-        "neg": "fisheye distortion, dark, muddy colours, ",
-    },
-    "cinematic": {
-        "keywords": ["cinematic", "movie", "scene", "dramatic", "epic",
-                     "action", "night", "rain", "fog", "smoke"],
-        "pos": "cinematic shot, anamorphic lens, film grain, color graded, "
-               "dramatic lighting, shallow depth of field, movie still, ",
-        "neg": "flat lighting, amateur, snapshot, overexposed, ",
-    },
-}
+# ── LLM prompt expansion ──────────────────────────────────────────────────────
+EXPANSION_SYSTEM = """You are an expert Stable Diffusion prompt engineer specialising in photorealistic and cinematic image generation.
+
+Your job: take a short user description and rewrite it as a detailed, accurate image generation prompt.
+
+Rules:
+- PRESERVE every specific detail from the input — if they say "one window open", "rainbow", "red door", those MUST appear
+- Wrap unique/specific details in attention weights like (one window open:1.4) or (rainbow:1.3)
+- Add: lighting description, camera/lens style, atmosphere, material textures, composition
+- Add quality boosters appropriate to the scene
+- Do NOT add people unless the user mentioned people
+- Do NOT change the subject or invent things not implied
+- Return ONLY the final prompt — no explanation, no preamble, no quotes
+- Keep it under 120 words
+- Use comma-separated tags and phrases, not full sentences"""
 
-# Things that are specific and easy to lose — boost their weight
-SPECIFIC_PATTERNS = [
-    r'\bone\b', r'\btwo\b', r'\bthree\b', r'\bsingle\b',
-    r'\bopen\b', r'\bclosed\b', r'\bbroken\b', r'\bempty\b',
-    r'\bright\b', r'\bleft\b', r'\btop\b', r'\bbottom\b',
-    r'rainbow', r'red ', r'blue ', r'green ', r'yellow ',
-    r'purple ', r'pink ', r'black ', r'white ', r'golden ',
-]
-
-def detect_scene(prompt_lower):
-    scores = {}
-    for scene, data in SCENE_TAGS.items():
-        score = sum(1 for kw in data["keywords"] if kw in prompt_lower)
-        if score > 0:
-            scores[scene] = score
-    if not scores:
-        return None
-    return max(scores, key=scores.get)
-
-def boost_specific_details(prompt):
-    """
-    Wraps specific/unique details in attention weights so the model
-    doesn't gloss over them. e.g. 'one window open' -> '(one window open:1.4)'
-    """
-    boosted = prompt
-
-    # Find phrases containing specific words and wrap them
-    specific_words = [
-        r'\b(one|single)\s+\w+(\s+\w+)?',      # "one window", "single door open"
-        r'\b(open|closed|broken|cracked)\s+\w+', # "open window", "broken glass"
-        r'\b\w+\s+(rainbow|aurora|lightning)\b',  # "rainbow over", "lightning bolt"
-        r'\b(rainbow|aurora|lightning)\b',
-        r'\b(left|right)\s+\w+',                 # "left side", "right hand"
-    ]
-
-    for pattern in specific_words:
-        def wrap(m):
-            return f"({m.group(0)}:1.4)"
-        boosted = re.sub(pattern, wrap, boosted, flags=re.IGNORECASE)
-
-    return boosted
-
-def expand_prompt(raw_prompt, style_choice):
-    """
-    Takes a short natural language prompt and expands it Fooocus-style.
-    Returns (expanded_positive, extra_negative)
-    """
-    prompt_lower = raw_prompt.lower()
-
-    # Detect scene
-    scene = detect_scene(prompt_lower)
-
-    # Start with quality prefix (added later outside this fn)
-    extra_pos = ""
-    extra_neg = ""
-
-    # Add scene vocabulary
-    if scene and style_choice == "Auto":
-        extra_pos += SCENE_TAGS[scene]["pos"]
-        extra_neg += SCENE_TAGS[scene]["neg"]
-
-    # Boost specific details in the original prompt
-    weighted_prompt = boost_specific_details(raw_prompt.strip())
-
-    # Add general realism boosters if no style override
-    if style_choice == "Auto":
-        extra_pos += "highly detailed, sharp focus, realistic, high resolution, "
-
-    return weighted_prompt, extra_pos, extra_neg
+def expand_prompt_llm(raw_prompt, style):
+    """Use Mistral to expand the user's short prompt Fooocus-style."""
+    if not raw_prompt.strip():
+        return ""
+
+    style_hint = f" The desired style is: {style}." if style != "Auto" else ""
+
+    user_msg = f"Expand this into a detailed image generation prompt:{style_hint}\n\n{raw_prompt.strip()}"
 
+    try:
+        response = llm_client.chat_completion(
+            messages=[
+                {"role": "system", "content": EXPANSION_SYSTEM},
+                {"role": "user",   "content": user_msg},
+            ],
+            max_tokens=200,
+            temperature=0.7,
+        )
+        expanded = response.choices[0].message.content.strip()
+        # Clean up any accidental quotes or preamble
+        expanded = expanded.strip('"').strip("'")
+        if expanded.lower().startswith("prompt:"):
+            expanded = expanded[7:].strip()
+        return expanded
+    except Exception as e:
+        print(f"LLM expansion failed, using raw prompt: {e}")
+        return raw_prompt.strip()
 
 # ── Style presets ─────────────────────────────────────────────────────────────
 STYLES = {
-    "Auto":        {"pos": "", "neg": ""},
-    "📸 Photo":    {
-        "pos": "RAW photo, photorealistic, DSLR, 8k uhd, film grain, Fujifilm XT3, ",
-        "neg": "painting, illustration, cartoon, anime, cgi, ",
+    "Auto":         {"pos": "", "neg": ""},
+    "📸 Photo":     {
+        "pos": "RAW photo, photorealistic, DSLR, 8k uhd, film grain, Fujifilm XT3, sharp focus, ",
+        "neg": "painting, illustration, cartoon, anime, cgi, render, ",
     },
     "🎬 Cinematic": {
-        "pos": "cinematic, movie still, anamorphic, film grain, color graded, dramatic lighting, ",
-        "neg": "flat, amateur, snapshot, overexposed, ",
+        "pos": "cinematic movie still, anamorphic lens, film grain, color graded, dramatic lighting, ",
+        "neg": "flat lighting, amateur, snapshot, overexposed, ",
     },
-    "🖼️ Portrait": {
-        "pos": "portrait, studio lighting, 85mm lens, bokeh, sharp eyes, detailed skin, ",
-        "neg": "wide angle, distorted, bad eyes, cropped, ",
+    "🖼️ Portrait":  {
+        "pos": "professional portrait, studio lighting, 85mm lens, bokeh, sharp eyes, skin texture, ",
+        "neg": "wide angle distortion, bad eyes, cropped head, ",
     },
     "🌆 Neon City": {
-        "pos": "cyberpunk, neon lights, rain reflections, night scene, blade runner aesthetic, ",
-        "neg": "daytime, rural, warm tones, ",
+        "pos": "cyberpunk city, neon lights, rain reflections, night scene, blade runner aesthetic, ",
+        "neg": "daytime, rural, nature, warm tones, ",
     },
-    "✨ Fantasy":   {
+    "✨ Fantasy":    {
         "pos": "fantasy art, epic, magical atmosphere, volumetric lighting, concept art, artstation, ",
         "neg": "modern, mundane, flat, ",
     },
-    "🎨 Painterly": {
+    "🎨 Painterly":  {
         "pos": "oil painting, impressionist, visible brushstrokes, canvas texture, museum quality, ",
-        "neg": "photo, digital art, flat colours, ",
+        "neg": "photo, digital flat art, ",
     },
 }
 
@@ -191,12 +133,13 @@ def generate(raw_prompt, negative_prompt, style, lora_name, lora_strength,
         seed = random.randint(0, 2**32 - 1)
     seed = int(seed)
 
-    # ── Smart expansion ──
-    weighted_prompt, scene_pos, scene_neg = expand_prompt(raw_prompt, style)
-    style_data = STYLES.get(style, STYLES["Auto"])
+    # ── LLM expansion ──
+    expanded = expand_prompt_llm(raw_prompt, style)
 
-    final_pos = IL_POS + style_data["pos"] + scene_pos + weighted_prompt
-    final_neg = IL_NEG + style_data["neg"] + scene_neg + negative_prompt.strip()
+    # ── Build final prompt ──
+    style_data = STYLES.get(style, STYLES["Auto"])
+    final_pos  = IL_POS + style_data["pos"] + expanded
+    final_neg  = IL_NEG + style_data["neg"] + negative_prompt.strip()
 
     # ── Move to GPU ──
     pipe.to("cuda")
@@ -206,7 +149,11 @@ def generate(raw_prompt, negative_prompt, style, lora_name, lora_strength,
     lora_data = LORAS.get(lora_name)
     if lora_data:
         try:
-            lp = hf_hub_download(repo_id=lora_data["repo"], filename=lora_data["file"])
+            lp = hf_hub_download(
+                repo_id=lora_data["repo"],
+                filename=lora_data["file"],
+                token=HF_TOKEN,
+            )
             pipe.load_lora_weights(lp)
             pipe.fuse_lora(lora_scale=float(lora_strength))
             lora_loaded = True
@@ -232,11 +179,10 @@ def generate(raw_prompt, negative_prompt, style, lora_name, lora_strength,
 
     pipe.to("cpu")
 
-    # Show expanded prompt if toggled on
-    expanded_text = f"**Sent to model:**\n{final_pos}" if show_expanded else ""
-
-    return result.images[0], seed, expanded_text
+    # ── Debug output ──
+    debug_text = f"**Expanded prompt sent to model:**\n\n{final_pos}" if show_expanded else ""
 
+    return result.images[0], seed, debug_text
 
 # ── CSS ───────────────────────────────────────────────────────────────────────
 css = """
@@ -250,7 +196,6 @@ body, .gradio-container {
     padding: 8px !important;
 }
 
-/* ── Topbar ── */
 .topbar {
     display: flex;
     align-items: center;
@@ -261,7 +206,6 @@ body, .gradio-container {
     color: #e8e0ff;
     font-size: 0.95em;
     font-weight: 800;
-    letter-spacing: -0.3px;
 }
 .gpu-pill {
     background: #1aff7a18;
@@ -275,7 +219,6 @@ body, .gradio-container {
     text-transform: uppercase;
 }
 
-/* ── Image output ── */
 .img-out {
     background: #0d0d1a;
     border: 1px solid #16162a;
@@ -283,7 +226,6 @@ body, .gradio-container {
     overflow: hidden;
     margin-bottom: 8px;
     min-height: 380px;
-    position: relative;
     display: flex;
     align-items: center;
     justify-content: center;
@@ -294,23 +236,16 @@ body, .gradio-container {
     display: block;
 }
 
-/* ── Seed pill under image ── */
-.seed-pill {
-    text-align: center;
-    margin-bottom: 12px;
-}
 .seed-pill input[type=number] {
     background: transparent !important;
     border: none !important;
     color: #2e2848 !important;
     font-size: 0.7em !important;
     text-align: center !important;
-    padding: 0 !important;
+    padding: 2px !important;
     width: 100% !important;
-    pointer-events: none;
 }
 
-/* ── Card ── */
 .card {
     background: #0d0d1a;
     border: 1px solid #16162a;
@@ -327,7 +262,6 @@ body, .gradio-container {
     margin-bottom: 8px;
 }
 
-/* ── Prompt textarea ── */
 textarea {
     background: transparent !important;
     border: none !important;
@@ -344,9 +278,9 @@ textarea::placeholder { color: #252038 !important; }
 textarea:focus {
     outline: none !important;
     box-shadow: none !important;
+    border: none !important;
 }
 
-/* ── Style pills ── */
 .style-wrap .gr-radio {
     display: flex !important;
     flex-wrap: wrap !important;
@@ -372,7 +306,6 @@ textarea:focus {
 }
 .style-wrap input[type=radio] { display: none !important; }
 
-/* ── Accordion ── */
 .gradio-accordion {
     background: #0d0d1a !important;
     border: 1px solid #16162a !important;
@@ -389,7 +322,6 @@ textarea:focus {
     padding: 12px 16px !important;
 }
 
-/* ── Sliders ── */
 .gradio-slider {
     background: transparent !important;
     border: none !important;
@@ -399,13 +331,7 @@ input[type=range] {
     accent-color: #6633bb !important;
     width: 100% !important;
 }
-.gradio-slider .wrap {
-    color: #6644aa !important;
-    font-size: 0.72em !important;
-    font-weight: 600 !important;
-}
 
-/* ── Number inputs ── */
 input[type=number] {
     background: #0a0a14 !important;
     border: 1px solid #18182a !important;
@@ -415,7 +341,6 @@ input[type=number] {
     padding: 8px 10px !important;
 }
 
-/* ── Checkbox ── */
 input[type=checkbox] { accent-color: #6633bb !important; }
 .gradio-checkbox label span {
     color: #4a3a6a !important;
@@ -423,37 +348,34 @@ input[type=checkbox] { accent-color: #6633bb !important; }
     font-weight: 600 !important;
 }
 
-/* ── Dropdown ── */
 .gradio-dropdown {
     background: #0a0a14 !important;
     border: 1px solid #18182a !important;
     border-radius: 10px !important;
 }
 
-/* ── Expanded prompt box ── */
-.expanded-box {
+label > span:first-child {
+    color: #3a2d55 !important;
+    font-size: 0.7em !important;
+    font-weight: 700 !important;
+    text-transform: uppercase !important;
+    letter-spacing: 1px !important;
+}
+
+/* Expanded prompt debug box */
+.debug-box {
     background: #080814;
     border: 1px solid #111122;
     border-radius: 10px;
     padding: 10px 12px;
-    color: #332255;
+    color: #443366;
     font-size: 0.7em;
-    line-height: 1.6;
+    line-height: 1.7;
     font-family: monospace;
     word-break: break-word;
-    min-height: 32px;
-}
-
-/* ── Labels ── */
-label > span:first-child {
-    color: #3a2d55 !important;
-    font-size: 0.7em !important;
-    font-weight: 700 !important;
-    text-transform: uppercase !important;
-    letter-spacing: 1px !important;
+    margin-bottom: 8px;
 }
 
-/* ── Generate button ── */
 .gen-btn button {
     background: linear-gradient(135deg, #4a1aaa 0%, #2d0e77 100%) !important;
     border: 1px solid #6633cc !important;
@@ -475,7 +397,6 @@ label > span:first-child {
 }
 .gen-btn button:active {
     transform: scale(0.98) !important;
-    box-shadow: 0 2px 12px #4a1aaa33 !important;
 }
 
 footer, .built-with { display: none !important; }
@@ -491,7 +412,6 @@ with gr.Blocks(css=css, title="ImageGen") as demo:
     </div>
     """)
 
-    # Output
     output_image = gr.Image(
         show_label=False, type="pil",
         height=460, elem_classes="img-out",
@@ -501,16 +421,14 @@ with gr.Blocks(css=css, title="ImageGen") as demo:
         elem_classes="seed-pill",
     )
 
-    # Prompt card
-    gr.HTML('<div class="card"><div class="card-label">Prompt</div>')
+    gr.HTML('<div class="card"><div class="card-label">✦ Prompt — write anything, short or long</div>')
     prompt = gr.Textbox(
         show_label=False,
-        placeholder="describe anything — short or long, it gets expanded automatically...",
+        placeholder="building with rainbow and one window open...",
         lines=3,
     )
     gr.HTML('</div>')
 
-    # Style pills
     gr.HTML('<div class="card-label" style="padding:4px 2px 8px;color:#3d3060;font-size:0.62em;font-weight:800;text-transform:uppercase;letter-spacing:2px;">Style</div>')
     style = gr.Radio(
         choices=list(STYLES.keys()),
@@ -519,13 +437,17 @@ with gr.Blocks(css=css, title="ImageGen") as demo:
         elem_classes="style-wrap",
     )
 
-    # Generate
     generate_btn = gr.Button(
         "Generate ✦", variant="primary",
         size="lg", elem_classes="gen-btn",
     )
 
-    # Advanced accordion
+    expanded_out = gr.Markdown(
+        value="",
+        elem_classes="debug-box",
+        visible=True,
+    )
+
     with gr.Accordion("⚙️  Settings", open=False):
         gr.HTML('<div style="height:6px"></div>')
 
@@ -542,7 +464,7 @@ with gr.Blocks(css=css, title="ImageGen") as demo:
             width  = gr.Slider(512, 1024, value=832,  step=64, label="Width")
             height = gr.Slider(512, 1216, value=1216, step=64, label="Height")
 
-        steps    = gr.Slider(20, 60,   value=30,  step=1,   label="Steps")
+        steps    = gr.Slider(20, 60,    value=30,  step=1,   label="Steps")
         guidance = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="CFG Scale")
 
         with gr.Row():
@@ -553,19 +475,15 @@ with gr.Blocks(css=css, title="ImageGen") as demo:
             randomize = gr.Checkbox(label="Random seed", value=True, scale=1)
 
         show_expanded = gr.Checkbox(
-            label="Show expanded prompt (debug)",
-            value=False,
+            label="Show expanded prompt (see what the LLM wrote)",
+            value=True,
         )
 
-    # LoRA accordion
     with gr.Accordion("🎨  LoRA", open=False):
         gr.HTML('<div style="height:6px"></div>')
         lora_name     = gr.Dropdown(choices=list(LORAS.keys()), value="None", label="LoRA preset")
         lora_strength = gr.Slider(0.1, 1.0, value=0.7, step=0.05, label="LoRA Strength")
 
-    # Expanded prompt debug output
-    expanded_out = gr.Markdown(elem_classes="expanded-box")
-
     generate_btn.click(
         fn=generate,
         inputs=[