adapt for llama 3.2

app.py

@@ -35,16 +35,22 @@ def infer_cap(image):
     return result
 
 
-model_path = "meta-llama/Llama-2-7b-chat-hf"
+import torch
+from transformers import pipeline
+
+llm_model_id = "meta-llama/Llama-3.2-3B-Instruct"
 
-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()
+llm_pipe = pipeline(
+        "text-generation",
+        model=llm_model_id,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+    )
 
 # 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
@@ -83,7 +89,7 @@ def infer_flux(prompt, seed=42, randomize_seed=True, width=1024, height=1024, gu
 @spaces.GPU
 def llama_gen_fragrance(scene):
 
-    instruction = """[INST] <<SYS>>\n
+    instruction = """
 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:
 ---
@@ -142,24 +148,27 @@ Always ensure that:
 – 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
+    messages = [
+        {"role": "system", "content": instruction},
+        {"role": "user", "content": scene},
+    ]
+    outputs = llm_pipe(
+        messages,
+        max_new_tokens=4096,
+    )
+    generated_response = outputs[0]["generated_text"][-1]['content']
+    print(generated_response)
+
+    return generated_response
 
 def extract_notes(text, section_name):
     import re
 
     # 1. Try block of bullets
-    pattern_block = rf'{section_name}:\s*\n((?:\*.*(?:\n|$))+)'
+    pattern_block = rf'{section_name}:\s*\n((?:\s*\*.*(?:\n|$))+)'
     match_block = re.search(pattern_block, text, re.MULTILINE)
     if match_block:
         notes_text = match_block.group(1)
@@ -168,11 +177,15 @@ def extract_notes(text, section_name):
             bullet = line.strip().lstrip('*').strip()
             if ':' in bullet:
                 note, desc = bullet.split(':', 1)
-                notes.append({'note': note.strip(), 'description': desc.strip()})
+            elif ',' in bullet:
+                note, desc = bullet.split(',', 1)
             else:
-                notes.append({'note': bullet, 'description': ''})
+                note, desc = bullet, ''
+            notes.append({'note': note.strip(), 'description': desc.strip()})
+
         return notes
 
+
     # 2. Try inline bullet style: * Section: item1, item2, item3
     pattern_inline = rf'\* {section_name}:\s*(.+)'
     match_inline = re.search(pattern_inline, text)
@@ -382,10 +395,10 @@ def infer(image_input):
 
     llama_q = moondream_result
     
-    gr.Info('Calling Llama2 ...')
+    gr.Info('Calling Llama3.2 ...')
     result = llama_gen_fragrance(llama_q)
 
-    print(f"Llama2 result: {result}")
+    #print(f"Llama2 result: {result}")
     yield result, None, None, None
 
     parsed = parse_perfume_description(result)
@@ -405,7 +418,7 @@ with gr.Blocks(css=css) as demo:
         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>
+            <p style="text-align: center">Upload an image, get a pro fragrance idea made by Llama !</p>
             """
         )
         with gr.Row():