""" Multimodal RAG Demo with Nemotron Embed VL and Rerank VL (ZeroGPU-friendly) Key ZeroGPU rule: - DO NOT load GPU models at import time. - Lazy-load models INSIDE the @spaces.GPU function (or inside helpers called from it). Models: - Embed: nvidia/llama-nemotron-embed-vl-1b-v2 - Rerank: nvidia/llama-nemotron-rerank-vl-1b-v2 - Gen (preferred): nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 (text-only summary, trust_remote_code) - If it fails, fallback to a smaller text-only model. Attention: - Default: SDPA (most stable on Spaces) - Optional: FlashAttention-2 if USE_FA2=1 and flash-attn is installed & compatible. """ import os import time import spaces import torch import gradio as gr from PIL import Image from datasets import load_dataset from safetensors.torch import load_file from transformers import ( AutoModel, AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer, AutoModelForCausalLM, ) # ----------------------------------------------------------------------------- # Config # ----------------------------------------------------------------------------- DEVICE_CPU = torch.device("cpu") EMBED_MODEL_PATH = "nvidia/llama-nemotron-embed-vl-1b-v2" EMBED_COMMIT_HASH = "5b5ca69c35bf6ec1484d2d5ff238626e67a745e2" RERANK_MODEL_PATH = "nvidia/llama-nemotron-rerank-vl-1b-v2" RERANK_COMMIT_HASH = "47e5a355d1a050c3e5f69d53f14964b1d34bcd9d" GENERATION_MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8" FALLBACK_GEN_MODEL_ID = os.getenv("FALLBACK_GEN_MODEL_ID", "Qwen/Qwen2.5-7B-Instruct") # ATTN_IMPL = "flash_attention_2" if os.getenv("USE_FA2", "0") == "1" else "sdpa" modality_to_tokens = {"image": 2048, "image_text": 10240, "text": 8192} PATH_TO_EMBEDDING_FILE = os.getenv("EMBEDDINGS_FILE", "image_text_embeddings_10k.safetensors") def check_flash_attention(): import torch from transformers.utils import is_flash_attn_2_available print(f"--- Flash Attention Check ---") print(f"PyTorch version: {torch.__version__}") print(f"CUDA available: {torch.cuda.is_available()}") # Transformers helper check fa2_available = is_flash_attn_2_available() print(f"Transformers reports FA2 available: {fa2_available}") if torch.cuda.is_available(): capability = torch.cuda.get_device_capability() print(f"GPU Compute Capability: {capability}") if capability[0] < 8: print("Note: FA2 requires Compute Capability 8.0+ (Ampere or newer).") return fa2_available # Determine best implementation if check_flash_attention(): ATTN_IMPL = "flash_attention_2" else: ATTN_IMPL = "sdpa" # Fallback to Scaled Dot Product Attention print(f"[INFO] Using {ATTN_IMPL} for model loading.") # model = AutoModelForCausalLM.from_pretrained( # model_id, # torch_dtype=torch.float16, # FA2 requires fp16 or bf16 # attn_implementation=best_attn, # trust_remote_code=True # ).to("cuda") # Call it inside your setup or first GPU call # FLASH_AVAILABLE = check_flash_attention() # ----------------------------------------------------------------------------- # Load dataset + embeddings (CPU only) # ----------------------------------------------------------------------------- print("[INFO] Loading dataset (CPU)...") dataset = load_dataset("mrdbourke/recipe-synthetic-images-10k") train_split = dataset["train"] print(f"[INFO] Dataset loaded with {len(train_split)} samples") # Pick the main markdown field robustly PREFERRED_TEXT_COL = "recipe_markdown" FALLBACK_TEXT_COLS = ["markdown", "text", "recipe_text", "content"] if PREFERRED_TEXT_COL in train_split.column_names: TEXT_COL = PREFERRED_TEXT_COL else: found = None for c in FALLBACK_TEXT_COLS: if c in train_split.column_names: found = c break if found is None: raise RuntimeError( f"Could not find a recipe text column. Available columns: {train_split.column_names}" ) TEXT_COL = found print(f"[WARN] '{PREFERRED_TEXT_COL}' not found. Using '{TEXT_COL}' instead.") if "image" not in train_split.column_names: raise RuntimeError(f"Dataset does not contain 'image' column. Columns: {train_split.column_names}") print(f"[INFO] Using TEXT_COL='{TEXT_COL}'") print("[INFO] Loading embeddings (CPU)...") emb = load_file(PATH_TO_EMBEDDING_FILE) if "image_text_embeddings" not in emb: raise RuntimeError(f"'{PATH_TO_EMBEDDING_FILE}' missing key 'image_text_embeddings'. Keys: {list(emb.keys())}") image_text_embeddings = emb["image_text_embeddings"].to(DEVICE_CPU) print(f"[INFO] Embeddings loaded: {tuple(image_text_embeddings.shape)} | device={image_text_embeddings.device}") # ----------------------------------------------------------------------------- # Lazy GPU globals (must only be initialized inside @spaces.GPU) # ----------------------------------------------------------------------------- _embed_model = None _embed_processor = None _rerank_model = None _rerank_processor = None _gen_model = None _gen_tokenizer = None _embeddings_gpu = None def _cuda() -> torch.device: return torch.device("cuda") def _load_embed_and_rerank_on_gpu(): global _embed_model, _embed_processor, _rerank_model, _rerank_processor device = _cuda() if _embed_model is None or _embed_processor is None: print("[INFO] Lazy-loading EMBED model on GPU...") _embed_model = AutoModel.from_pretrained( EMBED_MODEL_PATH, revision=EMBED_COMMIT_HASH, trust_remote_code=True, torch_dtype=torch.bfloat16, attn_implementation=ATTN_IMPL, ).to(device).eval() _embed_processor = AutoProcessor.from_pretrained( EMBED_MODEL_PATH, revision=EMBED_COMMIT_HASH, trust_remote_code=True, max_input_tiles=6, use_thumbnail=True, p_max_length=modality_to_tokens["image_text"], ) if _rerank_model is None or _rerank_processor is None: print("[INFO] Lazy-loading RERANK model on GPU...") _rerank_model = AutoModelForSequenceClassification.from_pretrained( RERANK_MODEL_PATH, revision=RERANK_COMMIT_HASH, trust_remote_code=True, torch_dtype=torch.bfloat16, attn_implementation=ATTN_IMPL, ).to(device).eval() _rerank_processor = AutoProcessor.from_pretrained( RERANK_MODEL_PATH, revision=RERANK_COMMIT_HASH, trust_remote_code=True, max_input_tiles=6, use_thumbnail=True, rerank_max_length=modality_to_tokens["image_text"], ) return _embed_model, _embed_processor, _rerank_model, _rerank_processor def _load_generation_model_on_gpu(): """ Try Nemotron 30B FP8 first. If it fails, fall back to a smaller text model. """ global _gen_model, _gen_tokenizer if _gen_model is not None and _gen_tokenizer is not None: return _gen_model, _gen_tokenizer device = _cuda() # 1) Try Nemotron FP8 try: print("[INFO] Lazy-loading GENERATION model (Nemotron 30B FP8) on GPU...") _gen_tokenizer = AutoTokenizer.from_pretrained( GENERATION_MODEL_ID, trust_remote_code=True, use_fast=True, ) # FIX: We force 'eager' or 'flash_attention_2' because this model # doesn't support the default 'sdpa' implementation yet. # Since you installed the FA2 wheels, we'll try to use that first. gen_attn_impl = ATTN_IMPL if ATTN_IMPL == "flash_attention_2" else "eager" _gen_model = AutoModelForCausalLM.from_pretrained( GENERATION_MODEL_ID, trust_remote_code=True, torch_dtype="auto", device_map="auto", attn_implementation=gen_attn_impl, # Changed here ).eval() print(f"[INFO] Nemotron generation model loaded OK with {gen_attn_impl}") return _gen_model, _gen_tokenizer except Exception as e: print(f"[WARN] Nemotron FP8 load failed: {repr(e)}") print(f"[WARN] Falling back to: {FALLBACK_GEN_MODEL_ID}") _gen_tokenizer = AutoTokenizer.from_pretrained( FALLBACK_GEN_MODEL_ID, trust_remote_code=True, use_fast=True, ) _gen_model = AutoModelForCausalLM.from_pretrained( FALLBACK_GEN_MODEL_ID, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto", attn_implementation="sdpa", # Fallback model usually supports SDPA ).eval() return _gen_model, _gen_tokenizer # ----------------------------------------------------------------------------- # Helpers # ----------------------------------------------------------------------------- def _l2_normalize(x: torch.Tensor, eps: float = 1e-12) -> torch.Tensor: return x / (x.norm(p=2, dim=-1, keepdim=True) + eps) def match_query_to_embeddings( query: str | Image.Image, target_embeddings_to_match: torch.Tensor, top_k: int = 50 ) -> tuple[torch.Tensor, torch.Tensor]: with torch.inference_mode(): if isinstance(query, Image.Image): q = _embed_model.encode_documents(images=[query]) else: q = _embed_model.encode_queries([query]) sim = _l2_normalize(q) @ _l2_normalize(target_embeddings_to_match).T sim = sim.flatten() idx = torch.argsort(sim, descending=True)[:top_k] scores = sim[idx] return scores, idx def rerank_samples( query_text: str, sorted_indices: torch.Tensor, num_samples_to_rerank: int = 20, ) -> tuple: device = _cuda() top_idx = sorted_indices[:num_samples_to_rerank] subset = dataset["train"].select(top_idx.tolist()) texts = subset[TEXT_COL] images = subset["image"] pairs = [{"question": query_text, "doc_text": t, "doc_image": im} for t, im in zip(texts, images)] batch = _rerank_processor.process_queries_documents_crossencoder(pairs) batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()} with torch.inference_mode(): out = _rerank_model(**batch, return_dict=True) logits = out.logits.squeeze(-1) rerank_sorted = torch.argsort(logits, descending=True) return subset, rerank_sorted def generate_recipe_summary(recipe_texts: list[str], max_new_tokens: int = 384) -> str: model, tok = _load_generation_model_on_gpu() combined = "" for i, r in enumerate(recipe_texts[:3], 1): combined += f"\n\n--- RECIPE {i} ---\n{r}" prompt = ( "You are a helpful culinary assistant.\n" "Summarize the following recipes in Markdown.\n\n" "Return:\n" "- 1–2 sentence overview of each\n" "- key ingredients\n" "- difficulty (Easy/Medium/Hard)\n" "- which is best for a quick weeknight dinner\n\n" f"{combined}\n\n" "## Summary:\n" ) inputs = tok(prompt, return_tensors="pt").to(model.device) with torch.inference_mode(): out = model.generate( **inputs, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.7, top_p=0.9, pad_token_id=tok.eos_token_id, ) gen = tok.decode(out[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True) return gen.strip() def _markdown_to_simple_html(markdown_text: str, max_reviews: int = 1) -> str: # Keep your original “card” parsing lightweight. lines = (markdown_text or "").strip().split("\n") title = "" description = "" cook_time = "" num_ratings = "" ingredients = [] steps = [] reviews = [] current_section = None in_ingredients = False in_steps = False in_reviews = False review_count = 0 for raw in lines: line = raw.strip() if line.startswith("# ") and not title: title = line[2:].strip() continue if line.startswith("**Time:**"): cook_time = line.replace("**Time:**", "").strip() continue if line.startswith("**Number of Ratings:**"): num_ratings = line.replace("**Number of Ratings:**", "").strip() continue if line.startswith("## "): section = line[3:].strip().lower() current_section = section in_ingredients = (section == "ingredients") in_steps = section.startswith("steps") in_reviews = (section == "reviews") continue if current_section == "description" and line and not line.startswith("#"): description = line continue if in_ingredients and line.startswith("- "): ingredients.append(line[2:].strip()) continue if in_steps and line and line[0].isdigit(): step_text = line.split(". ", 1)[-1] if ". " in line else line steps.append(step_text.strip()) continue if in_reviews and line.startswith("> ") and review_count < max_reviews: reviews.append(line[2:].strip()) review_count += 1 continue html = f"""
{title or "Recipe"}
{f"⏱️ {cook_time}" if cook_time else ""} {f"⭐ {num_ratings} ratings" if num_ratings else ""}
{(description[:160] + "…") if len(description) > 160 else description}
📝 Ingredients
{", ".join(ingredients[:8])}{("…" if len(ingredients) > 8 else "")}
👨‍🍳 Steps ({len(steps)})
    {"".join(f"
  1. {(s[:90] + '…') if len(s) > 90 else s}
    2. " for s in steps[:4])} {f"
  2. …and {len(steps)-4} more
    3. " if len(steps) > 4 else ""}
{f"
💬 Review
{(reviews[0][:220] + '…') if len(reviews[0]) > 220 else reviews[0]}
" if reviews else ""}
""" return html def create_recipe_cards_html(items: list[dict], num_results: int = 3) -> str: cards = [] for it in items[:num_results]: sample = it["sample"] md = sample.get(TEXT_COL, "") or "" cards.append(f"
{_markdown_to_simple_html(md)}
") return f"""

Retrieved Texts

{''.join(cards)}
""" # ----------------------------------------------------------------------------- # Main GPU function (ZeroGPU allocation happens here) # ----------------------------------------------------------------------------- @spaces.GPU def retrieve(query_text, query_image, rerank_option, generate_summary_option): global _embeddings_gpu # Load VL models only now _load_embed_and_rerank_on_gpu() if _embeddings_gpu is None: print("[INFO] Moving embeddings to GPU (cached)...") _embeddings_gpu = image_text_embeddings.to(_cuda(), non_blocking=True) # Choose query if query_text and str(query_text).strip(): input_query = str(query_text).strip() query_is_text = True elif query_image is not None: input_query = query_image query_is_text = False else: raise gr.Error("Please provide either a text query or an image query.") # Retrieval t0 = time.time() scores, idx = match_query_to_embeddings(input_query, _embeddings_gpu, top_k=20) t1 = time.time() top = dataset["train"].select(idx.tolist()) scored = [{"score": float(s.item()), "sample": smp} for s, smp in zip(scores, top)] gallery = [(it["sample"]["image"], f"Score: {it['score']:.4f}") for it in scored[:3]] cards_html = create_recipe_cards_html(scored, num_results=3) # Rerank (text only) if rerank_option == "True" and query_is_text: r0 = time.time() subset, rerank_sorted = rerank_samples(input_query, idx, num_samples_to_rerank=20) r1 = time.time() reranked = subset.select(rerank_sorted.tolist()) scored = [{"score": None, "sample": smp} for smp in reranked] gallery = [(it["sample"]["image"], f"Reranked: {i}") for i, it in enumerate(scored[:3])] cards_html = create_recipe_cards_html(scored, num_results=3) rerank_time = round(r1 - r0, 4) elif rerank_option == "True" and not query_is_text: rerank_time = "Reranking only supported for text queries" else: rerank_time = "Reranking turned off" # Generation (optional) if generate_summary_option == "True": g0 = time.time() recipe_texts = [it["sample"].get(TEXT_COL, "") for it in scored[:3]] summary = generate_recipe_summary(recipe_texts) summary = summary.replace("```markdown", "").replace("```", "").strip() g1 = time.time() gen_time = round(g1 - g0, 4) else: summary = "Generation turned off, no summary created" gen_time = "Generation turned off" timing = { "retrieve_time": round(t1 - t0, 4), "rerank_time": rerank_time, "generation_time": gen_time, "attn_impl": ATTN_IMPL, "text_col": TEXT_COL, } return gallery, cards_html, summary, timing # ----------------------------------------------------------------------------- # UI # ----------------------------------------------------------------------------- with gr.Blocks(title="Multimodal RAG Demo") as demo: gr.Markdown(f"""# 👁️📑 Multimodal RAG Demo (ZeroGPU-friendly) - Dataset: `mrdbourke/recipe-synthetic-images-10k` - Text field used: `{TEXT_COL}` - Embed: `nvidia/llama-nemotron-embed-vl-1b-v2` - Rerank: `nvidia/llama-nemotron-rerank-vl-1b-v2` - Gen (preferred): `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8` (text-only) - Attention backend: `{ATTN_IMPL}` (set `USE_FA2=1` to try FA2) """) with gr.Row(): with gr.Column(scale=1): query_text = gr.Textbox(label="Text Query", placeholder="e.g. 'dinner recipes with tomatoes'", lines=2) query_image = gr.Image(label="Image Query (optional)", type="pil", height=200) generate_summary_option = gr.Radio(["True", "False"], value="False", label="Generate recipe summary") rerank_option = gr.Radio(["True", "False"], value="False", label="Rerank initial results? (text only)") search_btn = gr.Button("Search", variant="primary") with gr.Column(scale=2): gallery_output = gr.Gallery(label="Retrieved Recipe Images", columns=3, height="auto", object_fit="cover") recipes_html = gr.HTML(label="Retrieved Recipe Texts") summary_generation = gr.Markdown(label="Generated Summary") timing_output = gr.JSON(label="Timings") gr.Examples( examples=[ ["best omelette recipes", None, "False", "False"], ["best omelette recipes", None, "False", "True"], ["eggplant dip", None, "True", "True"], ], inputs=[query_text, query_image, rerank_option, generate_summary_option], label="Example Queries", ) search_btn.click( fn=retrieve, inputs=[query_text, query_image, rerank_option, generate_summary_option], outputs=[gallery_output, recipes_html, summary_generation, timing_output], ) if __name__ == "__main__": demo.launch()