Upload 5 files

uno/utils/convert_yaml_to_args_file.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import yaml
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--yaml", type=str, required=True)
+parser.add_argument("--arg", type=str, required=True)
+args = parser.parse_args()
+
+
+with open(args.yaml, "r") as f:
+    data = yaml.safe_load(f)
+
+with open(args.arg, "w") as f:
+    for k, v in data.items():
+        if isinstance(v, list):
+            v = list(map(str, v))
+            v = " ".join(v)
+        if v is None:
+            continue
+        print(f"--{k} {v}", end=" ", file=f)

uno/utils/image_describer.py

@@ -0,0 +1,57 @@
+from transformers import Blip2Processor, Blip2ForConditionalGeneration
+from PIL import Image
+import torch
+import re
+
+BLIP2_MODEL_NAME = "Salesforce/blip2-flan-t5-xl"
+BLIP_DEVICE = "cpu"
+MAX_LENGTH = 120
+
+processor = None
+model = None
+
+def lazy_load_blip2():
+    global processor, model
+    if processor is None or model is None:
+        print("\U0001F680 [BLIP2] Loading BLIP-2 model and processor on CPU...")
+        processor = Blip2Processor.from_pretrained(BLIP2_MODEL_NAME)
+        model = Blip2ForConditionalGeneration.from_pretrained(
+            BLIP2_MODEL_NAME,
+            torch_dtype=torch.float32
+        ).to(BLIP_DEVICE).eval()
+
+def clean_caption(text: str) -> str:
+    text = text.strip()
+    text = re.sub(r"\s+", " ", text)
+    text = text.strip(' "\n')
+    return text[0].upper() + text[1:] if text else text
+
+def describe_uploaded_images(images: list[Image.Image]) -> dict:
+    if not images:
+        return {"style_description": "", "full_caption": ""}
+
+    lazy_load_blip2()
+
+    captions = []
+    prompt = (
+        "Describe this image in detail. Focus on the art medium, visual style, mood or tone, lighting or rendering cues, "
+        "and describe how people interact with objects if applicable."
+    )
+
+    for img in images:
+        try:
+            inputs = processor(images=img, text=prompt, return_tensors="pt").to(BLIP_DEVICE)
+            generated_ids = model.generate(**inputs, max_new_tokens=MAX_LENGTH)
+            caption = processor.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+            cleaned = clean_caption(caption)
+            if cleaned and cleaned not in captions:
+                captions.append(cleaned)
+        except Exception as e:
+            print(f"❌ [BLIP-2 ERROR] Failed to describe image: {e}")
+            continue
+
+    joined_caption = "; ".join(captions)
+    return {
+        "style_description": joined_caption,
+        "full_caption": joined_caption
+    }

uno/utils/prompt_enhancer.py

@@ -0,0 +1,150 @@
+import os
+import re
+from typing import List, Optional
+import openai
+
+from uno.utils.prompt_router import classify_prompt_intent
+from uno.utils.image_describer import describe_uploaded_images
+
+# === OpenAI Client ===
+client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
+# === Constants ===
+PROMPT_PREFIX = "It's very important that"
+PROMPT_SUFFIX = (
+    "and take all the time you needed as it is very important to achieve the best possible result"
+)
+TEMPLATE_FORMAT = "prefix [art medium] [main object or objective] [attribute] [expression] [key light] [detailing] suffix"
+DEFAULT_NUM_PROMPTS = 5
+MIN_WORD_COUNT = 8
+MAX_PROMPT_LENGTH = 120
+REQUIRED_ELEMENTS = [
+    "art medium",
+    "main object",
+    "attribute",
+    "expression",
+    "key light",
+    "detailing"
+]
+
+INTENT_SYSTEM_INSTRUCTIONS = {
+    "product_ad": "Focus: visually advertise a product using professional commercial photography.",
+    "service_promotion": "Focus: promote a service by illustrating its usage, setting, or emotional impact.",
+    "public_awareness": "Focus: support a cause, campaign, or message with narrative visual storytelling.",
+    "brand_storytelling": "Focus: express a brand's tone or identity using a visual lifestyle story.",
+    "creative_social_post": "Focus: generate stylistic or creative content suitable for social media that maintains core subject clarity.",
+    "fallback": "Fallback: default to showcasing the main product in a visually compelling commercial format."
+}
+
+def build_system_message(intent: str, num_prompts: int, style_hint: Optional[str] = "") -> str:
+    style_clause = (
+        f"The visual tone, lighting, and environment must match this style: {style_hint}.\n"
+        f"Only override this if the user explicitly requests a different visual style."
+        if style_hint else ""
+    )
+
+    if intent in INTENT_SYSTEM_INSTRUCTIONS:
+        instruction = INTENT_SYSTEM_INSTRUCTIONS[intent]
+    else:
+        print(f"🧠 [DEBUG] Unrecognized intent '{intent}', using dynamic fallback...")
+        instruction = (
+            f"Focus: Generate prompts suitable for a '{intent}' scenario using descriptive, high-quality visual storytelling. "
+            "The core subject must remain clear and central."
+        )
+
+    return f"""
+You are a prompt enhancement assistant for Flux Pro.
+
+Your task is to transform a short user input into {num_prompts} full-sentence, professional image generation prompts.
+
+Each prompt must follow this structure:
+{TEMPLATE_FORMAT}
+
+Prefix: '{PROMPT_PREFIX}'
+Suffix: '{PROMPT_SUFFIX}'
+
+Each prompt must:
+- Be under {MAX_PROMPT_LENGTH} words
+- Include: {", ".join(REQUIRED_ELEMENTS)}
+- Be a single descriptive sentence
+- Never use lists, examples, or bullet formatting
+- Avoid specific color names unless inferred from uploaded images
+- Do not wrap prompts in quotes or number them
+
+All image elements must follow natural physical proportions. Ensure objects intended for interaction appear in realistic size and position relative to the subject and scene. Avoid exaggerated scaling unless the user explicitly asks for surrealism or stylization.
+
+{style_clause}
+
+{instruction}
+
+Do not explain. Only return the prompts.
+Generate exactly {num_prompts} unique prompts, one per line.
+""".strip()
+
+
+def enhance_prompt_with_chatgpt(
+    user_prompt: str,
+    num_prompts: int = DEFAULT_NUM_PROMPTS,
+    reference_images: Optional[List] = None
+) -> List[str]:
+    intent = classify_prompt_intent(user_prompt)
+
+    blip_data = describe_uploaded_images(reference_images) if reference_images else {}
+    full_caption = blip_data.get("full_caption", "")
+    style_hint = blip_data.get("style_description", "")
+
+    print(f"\n📥 [DEBUG] User prompt: {user_prompt}")
+    if full_caption:
+        print(f"🖼️ [DEBUG] BLIP Caption: {full_caption}")
+    if style_hint:
+        print(f"🎨 [DEBUG] Style Description from Image: {style_hint}")
+
+    user_msg = f"Original prompt: {user_prompt}"
+    if style_hint:
+        user_msg += f"\nVisual reference style: {style_hint}"
+
+    try:
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[
+                {"role": "system", "content": build_system_message(intent, num_prompts, style_hint)},
+                {"role": "user", "content": user_msg}
+            ],
+            temperature=0.7,
+            max_tokens=1800
+        )
+
+        raw_output = response.choices[0].message.content.strip()
+        print("\n📝 [DEBUG] Raw GPT Output:")
+        print(raw_output)
+
+        candidate_prompts = [p.strip() for p in raw_output.split("\n") if p.strip()]
+
+        if (
+            len(candidate_prompts) == 1 and
+            candidate_prompts[0].lower().startswith(PROMPT_PREFIX.lower()) and
+            len(candidate_prompts[0].split()) > MIN_WORD_COUNT
+        ):
+            enhanced_prompts = [candidate_prompts[0]]
+        else:
+            enhanced_prompts = [
+                p for p in candidate_prompts
+                if len(p.split()) > MIN_WORD_COUNT and p.lower().startswith(PROMPT_PREFIX.lower())
+            ]
+
+        if len(enhanced_prompts) < num_prompts:
+            print(f"⚠️ Only {len(enhanced_prompts)} prompts returned. Padding with user prompt...")
+            enhanced_prompts += [user_prompt] * (num_prompts - len(enhanced_prompts))
+        elif len(enhanced_prompts) > num_prompts:
+            enhanced_prompts = enhanced_prompts[:num_prompts]
+
+        print("\n🔧 [DEBUG] ChatGPT Enhanced Prompts:")
+        for idx, p in enumerate(enhanced_prompts):
+            print(f"[{idx+1}] {p}")
+        print("--------------------------------------------------\n")
+
+        return enhanced_prompts
+
+    except Exception as e:
+        print(f"❌ [ERROR] Failed to enhance prompt: {e}")
+        return [user_prompt] * num_prompts

uno/utils/prompt_router.py

@@ -0,0 +1,48 @@
+import os
+import openai
+
+# === OpenAI Client ===
+client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
+# === System Message ===
+INTENT_SYSTEM_MESSAGE = """
+You are an intent classification engine for Flux Pro's AI prompt enhancement system.
+
+Given a short user prompt, your job is to classify the *underlying marketing or creative intent* behind it using a single lowercase label that reflects the intent category.
+
+You must:
+- Respond with only one concise intent label (a single word or hyphenated phrase)
+- Avoid using punctuation, explanations, or examples
+- Infer intent intelligently using your understanding of marketing and creative goals
+- You may return previously unseen or new intent labels if appropriate
+
+Some valid label types might include (but are not limited to): product-ad, service-promotion, public-awareness, brand-storytelling, social-trend, artistic-expression, educational-content, campaign-launch, or experimental-style.
+
+Only return the label. Do not echo the prompt or add any commentary.
+""".strip()
+
+def classify_prompt_intent(user_prompt: str) -> str:
+    """
+    Uses GPT-4 to classify the user's prompt into a marketing or creative intent label.
+    Returns a lowercase intent string (e.g., 'product-ad', 'artistic-expression').
+    """
+    user_msg = f"Prompt: {user_prompt}\nReturn only the label."
+
+    try:
+        response = client.chat.completions.create(
+            model="gpt-4",
+            messages=[
+                {"role": "system", "content": INTENT_SYSTEM_MESSAGE},
+                {"role": "user", "content": user_msg}
+            ],
+            temperature=0,
+            max_tokens=10,
+        )
+
+        label = response.choices[0].message.content.strip().lower()
+        print(f"📌 [DEBUG] Inferred intent label: {label}")
+        return label
+
+    except Exception as e:
+        print(f"❌ [ERROR] Failed to classify prompt intent: {e}")
+        return "unknown"