Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoProcessor, Llama4ForConditionalGeneration
+from transformers import AutoProcessor, AutoModelForCausalLM, AutoTokenizer
 import time
 import os
 from huggingface_hub import login
@@ -32,9 +32,8 @@ def test_text_generation():
             results.append(f"Error logging in: {e}")
             return "\n".join(results)
         
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-        
-        model_id = "meta-llama/Llama-4-8B-Instruct"  # Using smaller model for faster testing
+        # Use the correct model ID for Llama 4
+        model_id = "meta-llama/Llama-4-7b-instruct"  # Correct model ID
         
         results.append(f"Loading tokenizer from {model_id}...")
         tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
@@ -94,13 +93,14 @@ def test_image_text_generation():
             results.append(f"Error logging in: {e}")
             return "\n".join(results)
         
-        model_id = "meta-llama/Llama-4-Scout-8B-16E-Instruct"  # Using smaller model for faster testing
+        # Use the correct model ID for Llama 4 Scout
+        model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct"  # This should be correct
         
         results.append(f"Loading processor from {model_id}...")
         processor = AutoProcessor.from_pretrained(model_id, token=token)
         
         results.append(f"Loading model from {model_id}...")
-        model = Llama4ForConditionalGeneration.from_pretrained(
+        model = AutoModelForCausalLM.from_pretrained(
             model_id,
             token=token,
             torch_dtype=torch.bfloat16,
@@ -162,11 +162,66 @@ def test_image_text_generation():
     
     return "\n".join(results)
 
+# Function to list available Llama models
+def list_llama_models():
+    results = []
+    results.append("=== Listing Available Llama Models ===")
+    
+    try:
+        # Get token from environment
+        token = os.environ.get("HUGGINGFACE_TOKEN", "")
+        if token:
+            results.append(f"Token found: {token[:5]}...")
+        else:
+            results.append("No token found in environment variables!")
+            return "\n".join(results)
+        
+        # Login to Hugging Face
+        try:
+            login(token=token)
+            results.append("Successfully logged in to Hugging Face Hub")
+        except Exception as e:
+            results.append(f"Error logging in: {e}")
+            return "\n".join(results)
+        
+        # List models from the meta-llama organization
+        from huggingface_hub import HfApi
+        api = HfApi(token=token)
+        
+        results.append("Fetching models from meta-llama organization...")
+        models = api.list_models(author="meta-llama")
+        
+        results.append(f"Found {len(models)} models:")
+        for model in models:
+            results.append(f"- {model.id}")
+        
+        results.append("MODEL LISTING SUCCESSFUL!")
+        
+    except Exception as e:
+        results.append(f"Error listing models: {e}")
+        results.append(traceback.format_exc())
+        results.append("MODEL LISTING FAILED!")
+    
+    return "\n".join(results)
+
 # Create Gradio interface
 with gr.Blocks(title="Llama 4 Scout Test") as demo:
     gr.Markdown("# Llama 4 Scout Test")
     gr.Markdown("This Space tests the connection to Llama 4 and Llama 4 Scout models.")
     
+    with gr.Tab("List Available Models"):
+        with gr.Row():
+            with gr.Column():
+                list_models_button = gr.Button("List Available Llama Models")
+            with gr.Column():
+                models_result = gr.Textbox(label="Available Models", lines=20)
+        
+        list_models_button.click(
+            fn=list_llama_models,
+            inputs=[],
+            outputs=[models_result]
+        )
+    
     with gr.Tab("Text Generation Test"):
         with gr.Row():
             with gr.Column():
@@ -199,8 +254,9 @@ with gr.Blocks(title="Llama 4 Scout Test") as demo:
         
         This test checks if your Space can connect to and use Llama 4 and Llama 4 Scout models.
         
-        - The **Text Generation Test** uses the smaller Llama 4 8B model for basic text generation
-        - The **Image-Text Generation Test** uses the smaller Llama 4 Scout 8B model for image-text generation
+        - The **List Available Models** tab shows all models available from meta-llama
+        - The **Text Generation Test** uses Llama 4 7B for basic text generation
+        - The **Image-Text Generation Test** uses Llama 4 Scout for image-text generation
         
         If both tests pass, your Llama 4 Scout setup should work correctly.
         """)