Spaces:
Runtime error
Runtime error
| import torch | |
| import gradio as gr | |
| from PIL import Image | |
| from transformers import Blip2ForConditionalGeneration, AutoProcessor | |
| from peft import PeftModel, PeftConfig | |
| # Your fine-tuned model on Hugging Face | |
| peft_model_id = "MargiPandya/blip2-bdd100k" | |
| # Load PEFT configuration | |
| config = PeftConfig.from_pretrained(peft_model_id) | |
| # Load the base BLIP-2 model (use full precision if 8-bit not supported) | |
| model = Blip2ForConditionalGeneration.from_pretrained( | |
| config.base_model_name_or_path, | |
| torch_dtype=torch.float32, | |
| device_map=None, | |
| ) | |
| # Load the fine-tuned adapter | |
| model = PeftModel.from_pretrained(model, peft_model_id) | |
| # Load processor | |
| # processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b") | |
| processor = AutoProcessor.from_pretrained(config.base_model_name_or_path) | |
| # Move to CPU (for Spaces with no GPU) | |
| device = "cpu" | |
| model.to(device) | |
| model.eval() | |
| # Image captioning function | |
| def caption_image(image): | |
| if image is None: | |
| return "Please upload an image." | |
| # Ensure image is processed and input tensors are on same device as the model | |
| inputs = processor(images=image, text = "", return_tensors="pt") | |
| # Move all tensors to correct device | |
| inputs = {k: v.to(device) for k, v in inputs.items()} | |
| with torch.no_grad(): | |
| # Safe parameter recommendations for BLIP-2 | |
| generated_ids = model.generate( | |
| **inputs, | |
| max_new_tokens=100, | |
| do_sample=False, | |
| # temperature=0.9, | |
| # top_p=0.9, | |
| # repetition_penalty=1.5, | |
| ) | |
| caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] | |
| return caption.strip() | |
| # Gradio Interface | |
| app = gr.Interface( | |
| fn=caption_image, | |
| inputs=gr.Image(type="pil", label="Upload a driving scene image"), | |
| outputs=gr.Textbox(label="Generated Caption"), | |
| title="BLIP-2 BDD100K Scene Captioner", | |
| description="Upload a driving scene image and get an automatically generated caption.", | |
| theme="soft", | |
| ) | |
| if __name__ == "__main__": | |
| app.launch() | |