1Noura commited on
Commit
544f7f8
·
verified ·
1 Parent(s): 5e51ed4

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +54 -0
app.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ translator = pipeline(
15
+ task="translation",
16
+ model="facebook/nllb-200-distilled-600M",
17
+ torch_dtype=torch.bfloat16,
18
+ device=device
19
+ )
20
+
21
+ # Function to generate images based on the image's caption
22
+ def generate_image_and_translate(image, num_images=1):
23
+ caption_en = caption_image(image)[0]['generated_text']
24
+ caption_ar = translator(caption_en, src_lang="eng_Latn", tgt_lang="arb_Arab")[0]['translation_text']
25
+
26
+ generated_images = []
27
+ for _ in range(num_images):
28
+ generated_image = sd_pipeline(prompt=caption_en).images[0]
29
+ generated_images.append(generated_image)
30
+
31
+ return generated_images, caption_en, caption_ar
32
+
33
+ # Set up the Gradio interface
34
+ interface = gr.Interface(
35
+ fn=generate_image_and_translate,
36
+ inputs=[
37
+ gr.Image(type="pil", label="Upload Image"),
38
+ gr.Slider(minimum=1, maximum=10, label="Number of Images", value=1, step=1)
39
+ ],
40
+ outputs=[
41
+ gr.Gallery(label="Generated Images"),
42
+ gr.Textbox(label="Generated Caption (English)", interactive=False),
43
+ gr.Textbox(label="Translated Caption (Arabic)", interactive=False)
44
+ ],
45
+ title="Image Generation and Translation",
46
+ description="Upload an image to generate new images based on its caption and translate the caption into Arabic.",
47
+ examples=[
48
+ ["sea.jpg", 3]
49
+ ]
50
+ )
51
+
52
+ # Launch the Gradio application
53
+ if __name__ == "__main__":
54
+ interface.launch()