shahad-b commited on
Commit
dfe5b80
·
verified ·
1 Parent(s): 8963eca

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +72 -0
app.py ADDED
@@ -0,0 +1,72 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import pipeline
3
+ from diffusers import StableDiffusionPipeline
4
+ import torch
5
+ import wget
6
+
7
+ # Define the device to use (either "cuda" for GPU or "cpu" for CPU)
8
+ device = "cuda" if torch.cuda.is_available() else "cpu"
9
+
10
+ # Load the models
11
+ caption_image = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large", device=device)
12
+ sd_pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device)
13
+
14
+ # Load the translation model (English to Arabic)
15
+ translator = pipeline(
16
+ task="translation",
17
+ model="facebook/nllb-200-distilled-600M",
18
+ torch_dtype=torch.bfloat16,
19
+ device=device
20
+ )
21
+
22
+ # Download the image
23
+ url1 = "https://github.com/Shahad-b/Image-database/blob/main/sea.jpg?raw=true"
24
+ sea = wget.download(url1)
25
+
26
+ url2 = "https://github.com/Shahad-b/Image-database/blob/main/Cat.jpeg?raw=true"
27
+ Cat = wget.download(url2)
28
+
29
+ url3 = "https://github.com/Shahad-b/Image-database/blob/main/Car.jpeg?raw=true"
30
+ Car = wget.download(url3)
31
+
32
+ # Function to generate images based on the image's caption
33
+ def generate_image_and_translate(image, num_images=1):
34
+ # Generate caption in English from the uploaded image
35
+ caption_en = caption_image(image)[0]['generated_text']
36
+
37
+ # Translate the English caption to Arabic
38
+ caption_ar = translator(caption_en, src_lang="eng_Latn", tgt_lang="arb_Arab")[0]['translation_text']
39
+
40
+ generated_images = []
41
+
42
+ # Generate the specified number of images based on the English caption
43
+ for _ in range(num_images):
44
+ generated_image = sd_pipeline(prompt=caption_en).images[0]
45
+ generated_images.append(generated_image)
46
+
47
+ # Return the generated images along with both captions
48
+ return generated_images, caption_en, caption_ar
49
+
50
+ # Set up the Gradio interface
51
+ interface = gr.Interface(
52
+ fn=generate_image_and_translate,
53
+ inputs=[
54
+ gr.Image(type="pil", label="Upload Image"),
55
+ gr.Slider(minimum=1, maximum=10, label="Number of Images", value=1, step=1)
56
+ ],
57
+ outputs=[
58
+ gr.Gallery(label="Generated Images"),
59
+ gr.Textbox(label="Generated Caption (English)", interactive=False),
60
+ gr.Textbox(label="Translated Caption (Arabic)", interactive=False)
61
+ ],
62
+ title="Image Generation and Translation",
63
+ description="Upload an image to generate new images based on its caption and translate the caption into Arabic.",
64
+ examples=[
65
+ ["sea.jpg", 3],
66
+ ["Cat.jpeg", 4],
67
+ ["Car.jpeg", 2]
68
+ ]
69
+ )
70
+
71
+ # Launch the Gradio application
72
+ interface.launch()