app.py

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

from gradio_client import Client, handle_file

#client = Client("fffiloni/moondream2", hf_token=hf_token)

from transformers import AutoTokenizer, AutoModelForCausalLM

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


model_path = "meta-llama/Llama-2-7b-chat-hf"

tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False, use_auth_token=hf_token)
model = AutoModelForCausalLM.from_pretrained(model_path, use_auth_token=hf_token).half().cuda()

# FLUX

import numpy as np
import random
import torch
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(duration=75)
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):
- 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.
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.
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.**
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** section is always included and consistent with the rest.
– The Image Description must mention the exact Perfume Name on the label and focus exclusively on the perfume bottle as the main subject.
– Never mention or show humans, faces, body parts, characters, animals, or narrative props literally.
– 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)
    #print(generate_ids)
    #print(output_text)
    pattern = r'\[INST\].*?\[/INST\]'
    cleaned_text = re.sub(pattern, '', output_text, flags=re.DOTALL)
    return cleaned_text

def extract_notes(text, section_name):
    import re

    # 1. Try block of bullets
    pattern_block = rf'{section_name}:\s*\n((?:\*.*(?:\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():
            bullet = line.strip().lstrip('*').strip()
            if ':' in bullet:
                note, desc = bullet.split(':', 1)
                notes.append({'note': note.strip(), 'description': desc.strip()})
            else:
                notes.append({'note': bullet, 'description': ''})
        return notes

    # 2. 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 notes_raw.split(','):
            notes.append({'note': item.strip(), 'description': ''})
        return notes

    # 3. Try plain line style: Section: item1, item2, 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 notes_raw.split(','):
            notes.append({'note': item.strip(), 'description': ''})
        return notes

    return []

def parse_perfume_description(text: str) -> dict:
    # Perfume Name
    perfume_name = re.search(r'Perfume Name:\s*(.+)', text).group(1).strip()

    # Tagline (quoted)
    tagline = re.search(r'Tagline:\s*"(.*?)"', text, re.DOTALL)
    tagline = tagline.group(1).strip() if tagline else ""

    # Poetic Olfactory Description
    poetic_desc_match = re.search(
        r'Poetic Olfactory Description:\s*"(.*?)"', text, re.DOTALL)
    if poetic_desc_match:
        poetic_desc = poetic_desc_match.group(1).strip()
    else:
        poetic_desc_match = re.search(
            r'Poetic Olfactory Description:\s*(.*?)\s*(Olfactory Pyramid:|Image Description:|General Atmosphere:)', 
            text, re.DOTALL)
        poetic_desc = poetic_desc_match.group(1).strip() if poetic_desc_match else ""

    # General Atmosphere: stop at Image Description if present
    general_atmosphere_match = re.search(
        r'General Atmosphere:\s*(.*?)(?:\s*Image Description:|$)', text, re.DOTALL)
    general_atmosphere = general_atmosphere_match.group(1).strip() if general_atmosphere_match else ""

    # Image Description
    image_desc_match = re.search(
        r'Image Description:\s*"(.*?)"', text, re.DOTALL)
    if image_desc_match:
        image_desc = image_desc_match.group(1).strip()
    else:
        image_desc_match = re.search(
            r'Image Description:\s*(.*?)$', text, re.DOTALL)
        image_desc = image_desc_match.group(1).strip() if image_desc_match else ""

    # 🗂️ Smart bullet extractor
    
    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,
        '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()


def get_text_after_colon(input_text):
    # Find the first occurrence of ":"
    colon_index = input_text.find(":")
    
    # Check if ":" exists in the input_text
    if colon_index != -1:
        # Extract the text after the colon
        result_text = input_text[colon_index + 1:].strip()
        return result_text
    else:
        # Return the original text if ":" is not found
        return input_text

import pandas as pd

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

def extract_notes_for_comparison(data: Union[str, dict]) -> list[str]:
    """
    Extracts all notes from the Olfactory Pyramid section of a JSON string or dict.
    Args:
        data (Union[str, dict]): The JSON string or Python dict.
    Returns:
        list[str]: A list of extracted note names.
    """
    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")

    olfactory_pyramid = data.get("Olfactory Pyramid") or data.get("olfactory pyramid")
    if not olfactory_pyramid:
        raise KeyError("No 'Olfactory Pyramid' found in the data")

    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  # If a layer is missing, just skip
        for item in layer_data:
            note = item.get("note") or item.get("Note")
            if note:
                notes.append(note.strip())

    if not notes:
        raise ValueError("No notes found in the Olfactory Pyramid")

    return notes
    
from rapidfuzz import fuzz

def find_best_perfumes_from_json(data: Union[str, dict], top_n: int = 5, threshold: int = 80):
    """
    Finds top N matching perfumes using fuzzy matching on notes.
    Args:
        data (Union[str, dict]): The input JSON or dict.
        top_n (int): Number of results.
        threshold (int): Minimum fuzz ratio to count as match.
    Returns:
        pd.DataFrame
    """
    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:
                ratio = fuzz.partial_ratio(u_note, p_note)
                if ratio >= threshold:
                    matched_notes.append(p_note)

        matched_notes = sorted(set(matched_notes))

        matches.append({
            'brand': row['brand'],
            'perfume': row['perfume'],
            'matching_notes': ', '.join(matched_notes),
            'match_count': len(matched_notes)
        })

    result = pd.DataFrame(matches)
    result = result[result['match_count'] > 0]
    result = result.sort_values(by='match_count', 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)
    print(moondream_result)

    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)

    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]
    ).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)