Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import spaces
-from transformers import AutoModelForCausalLM, AutoProcessor
+from transformers import AutoModelForCausalLM, AutoProcessor, GPT2LMHeadModel, GPT2Tokenizer
 import torch
 from PIL import Image
 import subprocess
@@ -8,23 +8,45 @@ import subprocess
 # Install flash-attn with no CUDA build isolation
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-# Load model and processor
+# Define models and processors with pinning to a stable revision
 models = {
-    "microsoft/Phi-3.5-vision-instruct": AutoModelForCausalLM.from_pretrained("microsoft/Phi-3.5-vision-instruct", trust_remote_code=True, torch_dtype="auto", _attn_implementation="flash_attention_2").cuda().eval()
+    "microsoft/Phi-3.5-vision-instruct": AutoModelForCausalLM.from_pretrained(
+        "microsoft/Phi-3.5-vision-instruct",
+        revision="specific-revision-hash",  # Pinning to a specific revision for stability
+        trust_remote_code=True,
+        torch_dtype="auto", 
+        _attn_implementation="flash_attention_2"
+    ).cuda().eval()
 }
+
 processors = {
-    "microsoft/Phi-3.5-vision-instruct": AutoProcessor.from_pretrained("microsoft/Phi-3.5-vision-instruct", trust_remote_code=True)
+    "microsoft/Phi-3.5-vision-instruct": AutoProcessor.from_pretrained(
+        "microsoft/Phi-3.5-vision-instruct",
+        revision="specific-revision-hash",  # Pinning to a specific revision for stability
+        trust_remote_code=True
+    )
 }
 
+# Fallback to GPT-2 for testing
+def load_fallback_model():
+    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+    model = GPT2LMHeadModel.from_pretrained("gpt2").cuda().eval()
+    return model, tokenizer
+
 # Default description and prompt
-DESCRIPTION = ""
+DESCRIPTION = "[Phi-3.5-vision Demo](https://huggingface.co/microsoft/Phi-3.5-vision-instruct)"
 default_question = "You are an image to prompt converter. Your work is to observe each and every detail of the image and craft a detailed prompt under 100 words."
 
-# Gradio function for generating output from image input
+# Gradio function for generating output from image input with error handling
 @spaces.GPU
 def run_example(image, text_input=default_question, model_id="microsoft/Phi-3.5-vision-instruct"):
-    model = models[model_id]
-    processor = processors[model_id]
+    try:
+        model = models[model_id]
+        processor = processors[model_id]
+    except KeyError as e:
+        print(f"Error loading model: {e}. Falling back to GPT-2.")
+        model, processor = load_fallback_model()
+    
     user_prompt = '<|user|>\n'
     assistant_prompt = '<|assistant|>\n'
     prompt_suffix = "<|end|>\n"