llm

app.py

@@ -204,6 +204,23 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="blue", secondar
                 output_image = gr.Image(label="Result", elem_id="gallery", show_label=False)
                 used_seed = gr.Number(label="Seed Used")
 
+    def update_model_choices(provider):
+        provider_models = {
+            "Hugging Face": [
+                "Qwen/Qwen2.5-72B-Instruct",
+                "meta-llama/Meta-Llama-3.1-70B-Instruct",
+                "mistralai/Mixtral-8x7B-Instruct-v0.1",
+                "mistralai/Mistral-7B-Instruct-v0.3"
+            ],
+            "SambaNova": [
+                "Meta-Llama-3.1-70B-Instruct",
+                "Meta-Llama-3.1-405B-Instruct",
+                "Meta-Llama-3.1-8B-Instruct"
+            ],
+        }
+        models = provider_models.get(provider, [])
+        return gr.Dropdown(choices=models, value=models[0] if models else "")
+
     def update_llm_visibility(use_llm):
         return {
             llm_provider: gr.update(visible=use_llm),
@@ -217,15 +234,23 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="blue", secondar
         outputs=[llm_provider, llm_model, prompt_type]
     )
 
+    llm_provider.change(
+        update_model_choices,
+        inputs=[llm_provider],
+        outputs=[llm_model]
+    )
+
     def generate_prompt(image, text_prompt, use_enhancer, use_llm_generator, llm_provider, llm_model, prompt_type):
         if image is not None:
             caption = florence_caption(image)
-            if use_llm_generator:
-                prompt = generate_llm_prompt(caption, llm_provider, llm_model, prompt_type)
-            else:
-                prompt = caption
+            initial_prompt = caption
+        else:
+            initial_prompt = text_prompt
+        
+        if use_llm_generator:
+            prompt = generate_llm_prompt(initial_prompt, llm_provider, llm_model, prompt_type)
         else:
-            prompt = text_prompt
+            prompt = initial_prompt
         
         if use_enhancer:
             prompt = enhance_prompt(prompt)