Update app.py

app.py

@@ -2,13 +2,15 @@ import os
 from huggingface_hub import login
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
-import easyocr
+
 import gradio as gr
 from diffusers import DiffusionPipeline
 import torch
 import spaces  # Hugging Face Spaces module
 
-from transformers import pipeline
+from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
+
+
 
 
 
@@ -22,15 +24,15 @@ login(token=hf_token)
 processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
 model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
 
-
-pipe2= pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
+model2 = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 
 
-# Initialize the model
 pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 pipe.to(device)
-pipe2.to(device)
+model2.to(device)
 model.to(device)
 
 
@@ -41,6 +43,7 @@ def generate_caption_and_image(image):
     # reader = easyocr.Reader(['en'])  
     # result = reader.readtext(img)
     import random
+    
 
     # Define lists for the three variables
     fabrics = ['cotton', 'silk', 'denim', 'linen', 'polyester', 'wool', 'velvet']
@@ -54,15 +57,18 @@ def generate_caption_and_image(image):
     
     
 
-      
+    pixel_values = feature_extractor(images=[img], return_tensors="pt").pixel_values.to(device)
+    output_ids = model.generate(pixel_values, **gen_kwargs)
+    caption = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
+
 
     
     # Generate caption
     inputs = processor(image, return_tensors="pt", padding=True, truncation=True, max_length=250)
     inputs = {key: val.to(device) for key, val in inputs.items()}
     out = model.generate(**inputs)
-    caption = processor.decode(out[0], skip_special_tokens=True)
-    caption2 =pipe2(img)
+    caption2 = processor.decode(out[0], skip_special_tokens=True)
+
     prompt = f'''Create a highly realistic clothing item based on the following descriptions: The design should reflect {caption1} and {caption2}, blending both themes into a single, stylish, and modern piece of clothing. Incorporate highly realistic and high-quality textures that exude sophistication, with realistic fabric lighting and fine details. Subtly hint at {selected_fabric}, featuring a {selected_pattern} motif and a {selected_textile_design} style that harmoniously balances the essence of both captions.'''