Update app.py

app.py

@@ -21,6 +21,9 @@ model = AutoModelForCausalLM.from_pretrained(
 ).to(device)
 processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
 
+# Create a dummy image for text-only processing
+DUMMY_IMAGE = Image.new('RGB', (224, 224), color='white')
+
 # Load CivitAI dataset
 print("Loading dataset...")
 dataset = load_dataset("thefcraft/civitai-stable-diffusion-337k", split="train[:1000]")
@@ -31,18 +34,28 @@ text_embedding_cache = {}
 
 def get_image_embedding(image):
     try:
-        # Process image and add dummy text input
         inputs = processor(
             images=image,
-            text="Describe this image",  # Adding a default text prompt
-            padding=True,
-            return_tensors="pt"
+            text="Generate image description",
+            return_tensors="pt",
+            padding=True
         ).to(device, torch_dtype)
         
+        # Generate decoder_input_ids
+        decoder_input_ids = model.generate(
+            **inputs,
+            max_length=1,
+            min_length=1,
+            num_beams=1,
+            pad_token_id=processor.tokenizer.pad_token_id,
+            return_dict_in_generate=True,
+        ).sequences
+        
+        inputs['decoder_input_ids'] = decoder_input_ids
+        
         with torch.no_grad():
-            # Get model outputs
             outputs = model(**inputs)
-            # Extract image features from the cross-attention layers
+            # Use the mean of the last hidden state as the embedding
             image_embeddings = outputs.last_hidden_state.mean(dim=1)
         return image_embeddings.cpu().numpy()
     except Exception as e:
@@ -54,22 +67,26 @@ def get_text_embedding(text):
         if text in text_embedding_cache:
             return text_embedding_cache[text]
         
-        # Process text with proper input formatting
+        # Process text with dummy image
         inputs = processor(
+            images=DUMMY_IMAGE,
             text=text,
-            padding=True,
-            return_tensors="pt"
+            return_tensors="pt",
+            padding=True
         ).to(device, torch_dtype)
         
-        # Add required decoder input ids
-        inputs['decoder_input_ids'] = model.generate(
+        # Generate decoder_input_ids
+        decoder_input_ids = model.generate(
             **inputs,
             max_length=1,
+            min_length=1,
+            num_beams=1,
+            pad_token_id=processor.tokenizer.pad_token_id,
             return_dict_in_generate=True,
-            output_hidden_states=True,
-            early_stopping=True
         ).sequences
         
+        inputs['decoder_input_ids'] = decoder_input_ids
+        
         with torch.no_grad():
             outputs = model(**inputs)
             text_embeddings = outputs.last_hidden_state.mean(dim=1)
@@ -134,6 +151,9 @@ def process_image(input_image):
         if not isinstance(input_image, Image.Image):
             input_image = Image.fromarray(input_image)
         
+        # Resize image to expected size
+        input_image = input_image.resize((224, 224))
+        
         recommended_models, recommended_prompts = find_similar_images(input_image)
         
         if not recommended_models or not recommended_prompts: