mrm8488 commited on
Commit
ce0ce67
1 Parent(s): 4a462f4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -4
app.py CHANGED
@@ -1,13 +1,18 @@
1
  import torch
2
  from transformers import IdeficsForVisionText2Text, AutoProcessor
3
- from PIL import Image
4
  import gradio as gr
5
 
6
- model_id = "mrm8488/idefics-9b-ft-describe-diffusion-bf16"
7
  device = "cuda" if torch.cuda.is_available() else "cpu"
8
 
9
- model = IdeficsForVisionText2Text.from_pretrained(model_id, torch_dtype=torch.bfloat16)
10
- processor = AutoProcessor.from_pretrained(model_id)
 
 
 
 
 
11
 
12
  def predict(prompt, image_url, max_length):
13
  image = processor.image_processor.fetch_images(image_url)
 
1
  import torch
2
  from transformers import IdeficsForVisionText2Text, AutoProcessor
3
+ from peft import PeftModel, PeftConfig
4
  import gradio as gr
5
 
6
+ peft_model_id = "mrm8488/idefics-9b-ft-describe-diffusion-bf16-adapter"
7
  device = "cuda" if torch.cuda.is_available() else "cpu"
8
 
9
+
10
+ config = PeftConfig.from_pretrained(peft_model_id)
11
+ model = IdeficsForVisionText2Text.from_pretrained(config.base_model_name_or_path, torch_dtype=torch.bfloat16)
12
+ model = PeftModel.from_pretrained(model, peft_model_id)
13
+ processor = AutoProcessor.from_pretrained(config.base_model_name_or_path)
14
+ model = model.to(device)
15
+ model.eval()
16
 
17
  def predict(prompt, image_url, max_length):
18
  image = processor.image_processor.fetch_images(image_url)