sonu commited on
Commit
ab1e39e
·
1 Parent(s): a3c06ef

Add application file

Browse files
Files changed (3) hide show
  1. app.py +65 -0
  2. examples/example1.jpg +0 -0
  3. requirements.txt +5 -0
app.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from PIL import Image
4
+ from transformers import AutoProcessor, AutoModelForCausalLM
5
+
6
+ # Load model and processor
7
+ processor = AutoProcessor.from_pretrained("sonukiller/git-base-cartoon")
8
+ model = AutoModelForCausalLM.from_pretrained("sonukiller/git-base-cartoon")
9
+
10
+ # Move model to GPU if available
11
+ device = "cuda" if torch.cuda.is_available() else "cpu"
12
+ model = model.to(device)
13
+
14
+ def generate_caption(image):
15
+ """
16
+ Generate a caption for the given image using the custom model
17
+ """
18
+ # Preprocess the image
19
+ inputs = processor(images=image, return_tensors="pt").to(device)
20
+
21
+ # Generate caption
22
+ with torch.no_grad():
23
+ generated_ids = model.generate(
24
+ pixel_values=inputs.pixel_values,
25
+ max_length=50,
26
+ num_beams=4,
27
+ early_stopping=True
28
+ )
29
+
30
+ # Decode the generated ids to text
31
+ generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
32
+
33
+ return generated_caption
34
+
35
+ # Create Gradio interface
36
+ with gr.Blocks(title="Custom Image Captioning", css="footer {visibility: hidden}") as demo:
37
+ gr.Markdown("# Custom Image Captioning Model")
38
+ gr.Markdown("Upload an image and get a caption generated by a custom-trained model.")
39
+
40
+ with gr.Row():
41
+ with gr.Column():
42
+ input_image = gr.Image(type="pil", label="Input Image")
43
+ caption_button = gr.Button("Generate Caption")
44
+
45
+ with gr.Column():
46
+ output_text = gr.Textbox(label="Generated Caption")
47
+
48
+ caption_button.click(
49
+ fn=generate_caption,
50
+ inputs=input_image,
51
+ outputs=output_text
52
+ )
53
+
54
+ gr.Examples(
55
+ examples=[
56
+ "examples/example1.jpg",
57
+ ],
58
+ inputs=input_image,
59
+ outputs=output_text,
60
+ fn=generate_caption,
61
+ cache_examples=True,
62
+ )
63
+
64
+ # Launch the app
65
+ demo.launch()
examples/example1.jpg ADDED
requirements.txt ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ torch==2.0.1
2
+ Pillow==9.5.0
3
+ transformers==4.31.0
4
+ gradio==3.38.0
5
+ accelerate==0.21.0