Update app.py

app.py

@@ -6,45 +6,58 @@ import pandas as pd
 from datasets import load_dataset
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
+import warnings
+warnings.filterwarnings('ignore')
 
 # Load Florence-2 model and processor
 model_name = "microsoft/Florence-2-base"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
+# Modify model loading to disable flash attention
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype=torch_dtype,
-    trust_remote_code=True
+    trust_remote_code=True,
+    use_flash_attention=False  # Disable flash attention
 ).to(device)
 processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
 
 # Load CivitAI dataset (limited to 1000 samples)
+print("Loading dataset...")
 dataset = load_dataset("thefcraft/civitai-stable-diffusion-337k", split="train[:1000]")
 df = pd.DataFrame(dataset)
+print("Dataset loaded successfully!")
 
 # Create cache for embeddings to improve performance
 text_embedding_cache = {}
 
 def get_image_embedding(image):
-    inputs = processor(images=image, return_tensors="pt").to(device, torch_dtype)
-    with torch.no_grad():
-        outputs = model.get_image_features(**inputs)
-    return outputs.cpu().numpy()
+    try:
+        inputs = processor(images=image, return_tensors="pt").to(device, torch_dtype)
+        with torch.no_grad():
+            outputs = model.get_image_features(**inputs)
+        return outputs.cpu().numpy()
+    except Exception as e:
+        print(f"Error in get_image_embedding: {str(e)}")
+        return None
 
 def get_text_embedding(text):
-    if text in text_embedding_cache:
-        return text_embedding_cache[text]
-    
-    inputs = processor(text=text, return_tensors="pt").to(device, torch_dtype)
-    with torch.no_grad():
-        outputs = model.get_text_features(**inputs)
-    
-    embedding = outputs.cpu().numpy()
-    text_embedding_cache[text] = embedding
-    return embedding
+    try:
+        if text in text_embedding_cache:
+            return text_embedding_cache[text]
+        
+        inputs = processor(text=text, return_tensors="pt").to(device, torch_dtype)
+        with torch.no_grad():
+            outputs = model.get_text_features(**inputs)
+        
+        embedding = outputs.cpu().numpy()
+        text_embedding_cache[text] = embedding
+        return embedding
+    except Exception as e:
+        print(f"Error in get_text_embedding: {str(e)}")
+        return None
 
-# Pre-compute text embeddings for all prompts in the dataset
 def precompute_embeddings():
     print("Pre-computing text embeddings...")
     for idx, row in df.iterrows():
@@ -55,21 +68,21 @@ def precompute_embeddings():
     print("Finished pre-computing embeddings")
 
 def find_similar_images(uploaded_image, top_k=5):
-    # Get embedding for uploaded image
     query_embedding = get_image_embedding(uploaded_image)
+    if query_embedding is None:
+        return [], []
     
-    # Calculate similarities with dataset
     similarities = []
     for idx, row in df.iterrows():
         prompt_embedding = get_text_embedding(row['prompt'])
-        similarity = cosine_similarity(query_embedding, prompt_embedding)[0][0]
-        similarities.append({
-            'similarity': similarity,
-            'model': row['Model'],
-            'prompt': row['prompt']
-        })
+        if prompt_embedding is not None:
+            similarity = cosine_similarity(query_embedding, prompt_embedding)[0][0]
+            similarities.append({
+                'similarity': similarity,
+                'model': row['Model'],
+                'prompt': row['prompt']
+            })
     
-    # Sort by similarity and get top k results
     sorted_results = sorted(similarities, key=lambda x: x['similarity'], reverse=True)
     top_models = []
     top_prompts = []
@@ -94,21 +107,28 @@ def process_image(input_image):
     if input_image is None:
         return "Please upload an image.", "Please upload an image."
     
-    # Convert to PIL Image if needed
-    if not isinstance(input_image, Image.Image):
-        input_image = Image.fromarray(input_image)
-    
-    # Get recommendations
-    recommended_models, recommended_prompts = find_similar_images(input_image)
-    
-    # Format output
-    models_text = "Recommended Models:\n" + "\n".join([f"{i+1}. {model}" for i, model in enumerate(recommended_models)])
-    prompts_text = "Recommended Prompts:\n" + "\n".join([f"{i+1}. {prompt}" for i, prompt in enumerate(recommended_prompts)])
-    
-    return models_text, prompts_text
+    try:
+        if not isinstance(input_image, Image.Image):
+            input_image = Image.fromarray(input_image)
+        
+        recommended_models, recommended_prompts = find_similar_images(input_image)
+        
+        if not recommended_models or not recommended_prompts:
+            return "Error processing image.", "Error processing image."
+        
+        models_text = "Recommended Models:\n" + "\n".join([f"{i+1}. {model}" for i, model in enumerate(recommended_models)])
+        prompts_text = "Recommended Prompts:\n" + "\n".join([f"{i+1}. {prompt}" for i, prompt in enumerate(recommended_prompts)])
+        
+        return models_text, prompts_text
+    except Exception as e:
+        print(f"Error in process_image: {str(e)}")
+        return "Error processing image.", "Error processing image."
 
 # Pre-compute embeddings when starting the application
-precompute_embeddings()
+try:
+    precompute_embeddings()
+except Exception as e:
+    print(f"Error in precompute_embeddings: {str(e)}")
 
 # Create Gradio interface
 iface = gr.Interface(