app.py

import spaces
import gradio as gr
import numpy as np
import random
import torch
import re
import os 
import json
from typing import Union
hf_token = os.environ.get('HF_TOKEN')

from transformers import AutoTokenizer, AutoModelForCausalLM

# Load Moondream2 model for captioning
cap_model = AutoModelForCausalLM.from_pretrained(
    "vikhyatk/moondream2",
    revision="2025-06-21",
    trust_remote_code=True,
    device_map={"": "cuda"}  # ...or 'mps', on Apple Silicon
)

@spaces.GPU
def infer_cap(image):

    # Captioning
    #print("Short caption:")
    #print(model.caption(image, length="short")["caption"])
    
    cap = cap_model.caption(image, length="normal")["caption"]
    print("\nNormal caption:")
    print(cap)
    result = cap

    return result


# Load Llama LLM 

llm_model_path = "meta-llama/Llama-2-13b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(llm_model_path, use_fast=False, use_auth_token=hf_token)
model = AutoModelForCausalLM.from_pretrained(llm_model_path, use_auth_token=hf_token).half().cuda()

# FLUX

from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device)
torch.cuda.empty_cache()

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)

@spaces.GPU
def infer_flux(prompt, seed=42, randomize_seed=True, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    
    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
            prompt=prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
            output_type="pil",
            good_vae=good_vae,
        ):
            yield img
    
@spaces.GPU
def llama_gen_fragrance(scene):

    instruction = """[INST] <<SYS>>\n
You are a poetic perfumer. Your role is to create the imaginary scent of a described scene.
You must always respond using the following structure:
---
Perfume Name:
An original, evocative, and unique name — in French or English.
Tagline:
A short, poetic sentence — like a perfume advertisement hook.
Poetic Olfactory Description:
A freeform and expressive description of the scent ambiance evoked by the scene. Use rich sensory, emotional, and metaphorical language. Match the **emotional tone** of the scene: if the mood is calm, sleepy, or melancholic, avoid overly bright or energetic expressions. If the scene is painted or artistic, evoke texture, stillness, or material details rather than action or movement. Be subtle and precise.
**Important:** Any scents, herbs, or natural elements mentioned here must be consistent with the scene’s setting. Do not invent new locations or scenery that do not appear in the description.
Olfactory Pyramid (technical): 
**Important Consistency rule**: The top, heart, and base notes must not introduce new ideas, plants, or places that were not in the poetic description or the scene. Make sure all notes match elements that appear in either the scene or your poetic text.

Top Notes:
List 3–4 real, concrete scent materials that would be perceived first. These must be plausible fragrance ingredients (e.g. herbs, resins, citrus peels, spices, aldehydes, etc.). Pick notes that reflect the **real mood, climate, and setting** of the scene. Do not add locations or elements that don’t appear in the scene. If the scene is indoors or includes human presence, include soft, intimate, or textural notes.
Heart Notes:
List 3–4 real fragrance elements that give body and soul to the perfume. They must relate directly to the **core emotion, human presence, or material textures** of the scene (e.g. warm fabric, skin, dry flowers, books, wood, canvas). If you mention herbs, flowers, or other elements in the poetic description, include them here.
Base Notes:
List 3–4 real, longer-lasting ingredients such as woods, musks, resins, or earthy accords. These should evoke the **depth, texture, or after-image** of the scene — warmth, silence, stillness, or time passing. Avoid generic bases unless they fit the mood. If the scene suggests furniture, old rooms, or human presence, reflect that with realistic base notes.

General Atmosphere:
**This section is mandatory.** Provide a short, elegant paragraph summarizing the fragrance’s evolution and overall emotional impression. Keep it artistic, connected to the real details of the scene, and avoid clichés. **Never omit this section.**
Concentration Type:
**This section is mandatory.** After you write the General Atmosphere, decide which concentration best matches the scene’s mood and the perfume’s intended effect:
**Important:** Never default to Extrait de Parfum. Always justify your choice logically based on how strong or subtle the scent should feel for this scene, according to the Poetic Olfactory Description. The General Atmosphere must match and support the chosen concentration.
- **Extrait de Parfum (Parfum)**: Rich, deep, intense, very long-lasting — only choose this if the mood and atmosphere clearly require a powerful projection and deep longevity.
- **Eau de Parfum (EDP)**: Still rich and expressive, but more subtle than an Extrait. Perfect for warm, cozy, intimate, or nostalgic moods that don’t need overpowering strength.
- **Eau de Toilette (EDT)**: Lighter, fresher, good for daytime, breezy or casual scenes.
- **Eau Fraîche**: Very airy, fleeting, and refreshing — good for light, delicate, or spring-like moods.
- **Cologne**: Classic freshness, citrusy, herbal, and easy-going — evokes bright daytime or timeless lightness.
- **Body Mist / Brume Parfumée**: Softest and most subtle of all — ideal for gentle, barely-there scents, calming, dreamy, or bedroom-like atmospheres.
Image Description (for marketing visuals):
Describe an imagined marketing image that captures the perfume’s essence.
The Image Description **must** always begin by describing the perfume bottle as the clear, main, and visually dominant subject. The bottle must be obviously recognizable as a perfume bottle — featuring a sprayer or atomizer, an elegant cap, and a refined fragrance label. The label must clearly display the **exact Perfume Name** generated in this output, written exactly as it appears, along with a subtle, elegant mention of the imaginary luxury brand **“FILONI’S.”** The brand name should appear in a smaller, complementary font style, placed above or below the Perfume Name to enhance the overall design without overpowering it.
Do not use placeholder text like “Perfume Name” — always use the actual fragrance name exactly as you have suggested above.
**Important:** Absolutely never describe or depict any literal characters, humans, body parts, animals, narrative props, weapons, tools, furniture, costumes, or iconic objects from the input scene. Instead, translate any such elements into purely abstract or subtle design cues on the perfume bottle — for example, a hint of color, a texture, a minimal engraving, or an abstract shape. Never describe these objects directly. Never show them literally.
Describe the bottle’s shape, glass texture, cap, and label in fine detail. The glass may have an elegant frosted or matte finish, subtle etching or engraving (such as delicate floral or botanical motifs), or soft decorative elements that evoke refinement and sophistication while remaining tasteful and minimal.
Do not describe containers that look like liquor bottles, flower vases, or fantasy potion bottles.
Do not add narrative illustrations, characters, or storytelling scenes on the bottle — only subtle, abstract decorative details that highlight a luxury perfume aesthetic.
Specify the typography style used on the label text, ensuring it reflects the perfume’s mood and story (for example, elegant script for romantic scents, bold sans-serif for modern ones, vintage serif for nostalgic fragrances).
The bottle must occupy most of the image frame and appear in sharp focus and fine detail, shown from an angle or perspective that enhances its elegance and gives a refined, editorial feel — it does not have to be perfectly front-facing or centered.
Optionally, you may include one or two small, natural ingredients (such as herbs, flowers, citrus slices, or spices) placed tastefully near the bottle to subtly evoke the fragrance’s key notes — these must remain minimal and never overpower the bottle. 
**Only use ingredients that appear in the Olfactory Pyramid above** — do not invent or add any others. If you include an ingredient, depict it in a realistic, natural form.
The background should be minimal, abstract, or atmospheric — such as gradients, soft light, fabric textures, or mist — with no depiction of people, animals, or narrative scenes.
Use cinematic luxury advertising codes: refined shadows, soft directional lighting, elegant minimalism, and a sophisticated, editorial composition.
---
Always ensure that:
– The fragrance matches the mood and visual setting of the scene.
– All ingredients are real, plausible, and fit together naturally.
– No invented scenery or extra context is added.
– The poetic description and pyramid share the same notes and details.
– The **General Atmosphere** and **Concentration Type** sections are always included and consistent with the rest.
– The Concentration Type must always be justified by the General Atmosphere and mood from the Poetic Olfactory Description.
– The **Image Description** is always included and must mention the exact Perfume Name on the label and focus exclusively on the perfume bottle as the main subject.
– In Image Description, Never mention or show humans, faces, body parts, characters, animals, or narrative props literally.
– In Image Description, Any props, costumes, or iconic objects must be abstracted into subtle decorative or textural cues only.
– Never describe these narrative elements directly.
– Each perfume feels unique and consistent.
Here is the scene description to analyze:
\n<</SYS>>\n\n{} [/INST]"""
   
    prompt = instruction.format(scene)
    
    generate_ids = model.generate(tokenizer(prompt, return_tensors='pt').input_ids.cuda(), max_new_tokens=4096)
    output_text = tokenizer.decode(generate_ids[0], skip_special_tokens=True)

    pattern = r'\[INST\].*?\[/INST\]'
    cleaned_text = re.sub(pattern, '', output_text, flags=re.DOTALL)
    return cleaned_text

def extract_notes(text, section_name):

    bullet_pattern = r'(\*|\d+[a-zA-Z]?\.|-|•|\+|[a-zA-Z]\.)'

    # 1. Try block of bullets with all bullet types
    pattern_block = rf'{section_name}:\s*\n((?:\s*{bullet_pattern}\s*.*(?:\n|$))+)'
    match_block = re.search(pattern_block, text, re.MULTILINE)
    if match_block:
        notes_text = match_block.group(1)
        notes = []
        for line in notes_text.strip().splitlines():
            line = line.strip()
            bullet = re.sub(rf'^{bullet_pattern}\s*', '', line)
            if ':' in bullet:
                note, desc = bullet.split(':', 1)
            elif ',' in bullet:
                note, desc = bullet.split(',', 1)
            else:
                note, desc = bullet, ''
            notes.append({'note': note.strip(), 'description': desc.strip()})
        return notes

    # 2. Try multiline block without bullets (just lines with commas/and)
    pattern_multiline = rf'{section_name}:\s*\n((?:[^\n]*\n)+)'
    match_multiline = re.search(pattern_multiline, text, re.MULTILINE)
    if match_multiline:
        # Get all lines after the section header until next empty line or another section
        block = match_multiline.group(1).strip()
        # Stop at next section header if any
        block = re.split(r'\n\S+:\s*\n', block)[0].strip()
        # Now split by commas or 'and'
        notes = []
        for item in re.split(r',|\band\b', block):
            item = item.strip().strip('.')
            if item:
                notes.append({'note': item, 'description': ''})
        if notes:
            return notes

    # 3. Try inline bullet style: * Section: item1, item2, item3
    pattern_inline = rf'\* {section_name}:\s*(.+)'
    match_inline = re.search(pattern_inline, text)
    if match_inline:
        notes_raw = match_inline.group(1).strip()
        notes = []
        for item in re.split(r',|\band\b', notes_raw):
            item = item.strip().strip('.')
            if item:
                notes.append({'note': item, 'description': ''})
        return notes

    # 4. Try plain line style: Section: item1, item2, and item3 (no bullet)
    pattern_line = rf'^{section_name}:\s*(.+)$'
    match_line = re.search(pattern_line, text, re.MULTILINE)
    if match_line:
        notes_raw = match_line.group(1).strip()
        notes = []
        for item in re.split(r',|\band\b', notes_raw):
            item = item.strip().strip('.')
            if item:
                notes.append({'note': item, 'description': ''})
        return notes

    return []

def parse_perfume_description(text: str) -> dict:
    def safe_search(pattern, text, flags=0):
        match = re.search(pattern, text, flags)
        return match.group(1).strip() if match else None

    # Perfume Name
    perfume_name = safe_search(r'Perfume Name:\s*(.+)', text) or ""

    # Concentration Type (same style: expects `Concentration Type:`)
    concentration_type = safe_search(r'Concentration Type:\s*(.+)', text) or ""

    # Tagline
    tagline = safe_search(r'Tagline:\s*"(.*?)"', text, re.DOTALL) or ""

    # Poetic Olfactory Description
    poetic_desc = safe_search(r'Poetic Olfactory Description:\s*"(.*?)"', text, re.DOTALL)
    if poetic_desc is None:
        poetic_desc = safe_search(
            r'Poetic Olfactory Description:\s*(.*?)\s*(Olfactory Pyramid:|Image Description:|General Atmosphere:|Concentration Type:)', 
            text, re.DOTALL) or ""

    # General Atmosphere
    general_atmosphere = safe_search(
        r'General Atmosphere:\s*(.*?)(?:\s*Image Description:|$)', 
        text, re.DOTALL) or ""

    # Image Description
    image_desc = safe_search(r'Image Description:\s*"(.*?)"', text, re.DOTALL)
    if image_desc is None:
        image_desc = safe_search(r'Image Description:\s*(.*?)$', text, re.DOTALL) or ""

    # Notes
    top_notes = extract_notes(text, 'Top Notes')
    heart_notes = extract_notes(text, 'Heart Notes')
    base_notes = extract_notes(text, 'Base Notes')

    result = {
        'Perfume Name': perfume_name,
        'Concentration Type': concentration_type,
        'Tagline': tagline,
        'Poetic Olfactory Description': poetic_desc,
        'Image Description': image_desc,
        'Olfactory Pyramid': {
            'Top Notes': top_notes,
            'Heart Notes': heart_notes,
            'Base Notes': base_notes
        },
        'General Atmosphere': general_atmosphere
    }

    return result

def extract_field(data: Union[str, dict], field_name: str) -> str:
    """
    Extracts a specific field value from a JSON string or Python dict.
    Args:
        data (Union[str, dict]): The JSON string or dict to extract from.
        field_name (str): The exact field name to extract.
    Returns:
        str: The extracted field value as a string.
    """
    if isinstance(data, str):
        try:
            data = json.loads(data)
        except json.JSONDecodeError:
            raise ValueError("Invalid JSON string provided")

    if not isinstance(data, dict):
        raise TypeError("Input must be a dict or a valid JSON string")

    value = data.get(field_name) or data.get(field_name.lower()) or None

    if value is None:
        raise KeyError(f"No field named '{field_name}' found in the data")

    return str(value).strip()


import pandas as pd
from rapidfuzz import fuzz

# Load the database once
df = pd.read_excel('perfume_database_cleaned.xlsx')

def extract_notes_for_comparison(data: Union[str, dict]) -> list[str]:
    """
    Extracts notes from Olfactory Pyramid in a JSON string or dict.
    """
    if isinstance(data, str):
        try:
            data = json.loads(data)
        except json.JSONDecodeError:
            raise ValueError("Invalid JSON string provided")

    if not isinstance(data, dict):
        raise TypeError("Input must be a dict or valid JSON string")

    olfactory_pyramid = data.get("Olfactory Pyramid") or data.get("olfactory pyramid")
    if not olfactory_pyramid:
        return []  # No pyramid found, fail gracefully

    notes = []
    for layer in ["Top Notes", "Heart Notes", "Base Notes"]:
        layer_data = olfactory_pyramid.get(layer) or olfactory_pyramid.get(layer.lower())
        if not layer_data:
            continue
        for item in layer_data:
            note = item.get("note") or item.get("Note")
            if note:
                notes.append(note.strip())

    return notes

def find_best_perfumes_from_json(data: Union[str, dict], top_n: int = 5, threshold: int = 80):
    """
    Fuzzy-match user notes against database notes.
    Uses token_set_ratio + partial_ratio + short-word safeguard.
    """
    user_notes = extract_notes_for_comparison(data)
    user_notes_clean = [n.strip().lower() for n in user_notes]

    matches = []

    for _, row in df.iterrows():
        perfume_notes = [n.strip().lower() for n in row['notes'].split(',')]
        matched_notes = []

        for u_note in user_notes_clean:
            for p_note in perfume_notes:
                if len(u_note) < 4:
                    # Very short? Require exact match
                    ratio = 100 if u_note == p_note else 0
                else:
                    ratio_token = fuzz.token_set_ratio(u_note, p_note)
                    ratio_partial = fuzz.partial_ratio(u_note, p_note)
                    ratio = max(ratio_token, ratio_partial)

                if ratio >= threshold:
                    matched_notes.append(p_note)

        unique_matched_notes = sorted(set(matched_notes))
        match_count = len(unique_matched_notes)

        if match_count == 0:
            continue  # Skip if no match at all

        total_notes = len(perfume_notes)
        purity = match_count / total_notes if total_notes > 0 else 0
        adjusted_score = match_count * purity

        matches.append({
            'brand': row['brand'],
            'perfume': row['perfume'],
            'matching_notes': ', '.join(unique_matched_notes).strip(', '),
            'match_count': match_count,
            'purity': round(purity, 2),
            'adjusted_score': round(adjusted_score, 2)
        })

    if not matches:
        # Nothing matched at all
        return pd.DataFrame([{
            'brand': 'N/A',
            'perfume': 'No match found',
            'matching_notes': '',
            'match_count': 0,
            'purity': 0,
            'adjusted_score': 0
        }])

    result = pd.DataFrame(matches)
    result = result.sort_values(by='adjusted_score', ascending=False).head(top_n).reset_index(drop=True)

    return result
    
def infer(image_input):

    gr.Info('Calling Moondream model for caption ...')
    yield None, None, None, None
    moondream_result = infer_cap(image_input)

    llama_q = moondream_result
    
    gr.Info('Calling Llama2 ...')
    result = llama_gen_fragrance(llama_q)

    print(f"Llama2 result: {result}")
    yield result, None, None, None

    parsed = parse_perfume_description(result)
    image_desc = extract_field(parsed, "Image Description")

    real_correspondance = find_best_perfumes_from_json(parsed)
    
    yield result, parsed, image_desc, real_correspondance

css="""
#col-container {max-width: 910px; margin-left: auto; margin-right: auto;}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            <h1 style="text-align: center">Image to Fragrance</h1>
            <p style="text-align: center">Upload an image, get a pro fragrance idea made by Llama2 !</p>
            """
        )
        with gr.Row():
            with gr.Column():
                image_in = gr.Image(label="Image input", type="pil", elem_id="image-in")
                submit_btn = gr.Button('Give me a Fragrance')
                json_res = gr.JSON(label="JSON (for further usage)")
                flacon_desc = gr.Textbox(interactive=False, visible=False)
            with gr.Column():
                #caption = gr.Textbox(label="Generated Caption")
                fragrance = gr.Textbox(label="generated Fragrance", elem_id="fragrance")
                output_df = gr.Dataframe(visible=False)
                get_flacon_btn = gr.Button("Generate Flacon image", interactive=False)
                bottle_res = gr.Image(label="Flacon")

    def disable_flacon_button():         
        return gr.update(interactive=False), gr.update(visible=False), gr.update(value=None)

    def allow_flacon_button():         
        return gr.update(interactive=True), gr.update(visible=True)
    
    submit_btn.click(
        fn=disable_flacon_button,
        inputs = [],
        outputs = [get_flacon_btn, output_df, bottle_res]
    ).then(
        fn=infer, 
        inputs=[image_in], 
        outputs=[fragrance, json_res, flacon_desc, output_df]
    ).then(
        fn=allow_flacon_button,
        inputs=[],
        outputs=[get_flacon_btn, output_df]
    )
    get_flacon_btn.click(fn=infer_flux, inputs=[flacon_desc], outputs=[bottle_res])

demo.queue(max_size=12).launch(ssr_mode=False, mcp_server=True)