Use direct nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4 model API in generate_prompt

app.py

@@ -2,6 +2,7 @@ import os
 import json
 import gradio as gr
 from gradio_client import Client
+from huggingface_hub import InferenceClient
 
 
 def generate_prompt(concept: str) -> str:
@@ -11,7 +12,7 @@ def generate_prompt(concept: str) -> str:
     if not concept:
         return "a ginger cat wearing a tiny wizard hat reading a spellbook"
     try:
-        client = Client("akhaliq/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4", verbose=False)
+        client = InferenceClient("nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4")
         system_instruction = (
             "You are an expert prompt engineer for text-to-image models. "
             "Your task is to take a simple concept and expand it into a detailed, "
@@ -24,13 +25,12 @@ def generate_prompt(concept: str) -> str:
             {"role": "system", "content": system_instruction},
             {"role": "user", "content": f"Concept: {concept}"}
         ]
-        result = client.predict(
-            messages_json=json.dumps(messages),
+        response = client.chat_completion(
+            messages=messages,
             temperature=0.7,
-            max_tokens=256,
-            custom_token=None,
-            api_name="/chat"
+            max_tokens=256
         )
+        result = response.choices[0].message.content
         clean_result = str(result).strip()
         if clean_result.startswith('"') and clean_result.endswith('"'):
             clean_result = clean_result[1:-1]
@@ -38,7 +38,7 @@ def generate_prompt(concept: str) -> str:
             clean_result = clean_result[1:-1]
         return clean_result
     except Exception as e:
-        print(f"Error calling Nemotron space: {e}")
+        print(f"Error calling Nemotron model: {e}")
         return f"A detailed, high-quality, professional commercial product photograph of {concept}"