gokilashree commited on
Commit
9f1b2f8
·
verified ·
1 Parent(s): 2c67ddf

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +98 -0
app.py ADDED
@@ -0,0 +1,98 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import MBartForConditionalGeneration, MBart50Tokenizer, AutoModelForCausalLM, AutoTokenizer, pipeline
3
+ from diffusers import DiffusionPipeline
4
+ import torch
5
+ from PIL import Image
6
+
7
+ # Load the Translation Model (MBART for Tamil to English Translation)
8
+ model_name = "facebook/mbart-large-50-many-to-one-mmt"
9
+ tokenizer = MBart50Tokenizer.from_pretrained(model_name)
10
+ model = MBartForConditionalGeneration.from_pretrained(model_name)
11
+
12
+ # Load the Text Generation Model (for generating a short paragraph)
13
+ text_generation_model_name = "EleutherAI/gpt-neo-1.3B"
14
+ text_tokenizer = AutoTokenizer.from_pretrained(text_generation_model_name)
15
+ text_model = AutoModelForCausalLM.from_pretrained(text_generation_model_name)
16
+ text_generator = pipeline("text-generation", model=text_model, tokenizer=text_tokenizer)
17
+
18
+ # Load the Stable Diffusion XL Model for Image Generation
19
+ pipe = DiffusionPipeline.from_pretrained(
20
+ "stabilityai/stable-diffusion-xl-base-1.0",
21
+ torch_dtype=torch.float16,
22
+ variant="fp16",
23
+ )
24
+ pipe.to("cuda") # Send the model to GPU for faster image generation
25
+
26
+ # Function to generate image from text prompt using Stable Diffusion XL
27
+ def generate_image_from_text(translated_text):
28
+ try:
29
+ print(f"Generating image from translated text: {translated_text}")
30
+ # Generate the image using the pipeline
31
+ image = pipe(prompt=translated_text).images[0]
32
+ print("Image generation completed.")
33
+ return image
34
+ except Exception as e:
35
+ print(f"Error during image generation: {e}")
36
+ return None
37
+
38
+ # Function to generate a short paragraph from the translated text
39
+ def generate_short_paragraph_from_text(translated_text):
40
+ try:
41
+ print(f"Generating a short paragraph from translated text: {translated_text}")
42
+ paragraph = text_generator(
43
+ translated_text,
44
+ max_length=80, # Reduced to 80 tokens
45
+ num_return_sequences=1,
46
+ temperature=0.6,
47
+ top_p=0.8,
48
+ truncation=True # Added truncation to avoid long sequences
49
+ )[0]['generated_text']
50
+ print(f"Paragraph generation completed: {paragraph}")
51
+ return paragraph
52
+ except Exception as e:
53
+ print(f"Error during paragraph generation: {e}")
54
+ return f"Error during paragraph generation: {e}"
55
+
56
+ # Function to translate Tamil text, generate a short paragraph, and create an image
57
+ def translate_generate_paragraph_and_image(tamil_text):
58
+ # Step 1: Translate Tamil text to English using mbart-large-50
59
+ try:
60
+ print("Translating Tamil text to English...")
61
+ tokenizer.src_lang = "ta_IN"
62
+ inputs = tokenizer(tamil_text, return_tensors="pt")
63
+ translated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id["en_XX"])
64
+ translated_text = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
65
+ print(f"Translation completed: {translated_text}")
66
+ except Exception as e:
67
+ return f"Error during translation: {e}", "", None
68
+
69
+ # Step 2: Generate a shorter paragraph based on the translated English text
70
+ paragraph = generate_short_paragraph_from_text(translated_text)
71
+ if "Error" in paragraph:
72
+ return translated_text, paragraph, None
73
+
74
+ # Step 3: Generate an image using the translated English text with the new model
75
+ image = generate_image_from_text(translated_text)
76
+
77
+ return translated_text, paragraph, image
78
+
79
+ # Define Gradio Interface
80
+ def interface(tamil_text):
81
+ translated_text, paragraph, image = translate_generate_paragraph_and_image(tamil_text)
82
+ return translated_text, paragraph, image
83
+
84
+ # Create Gradio Interface (with the image output)
85
+ iface = gr.Interface(
86
+ fn=interface,
87
+ inputs=gr.Textbox(lines=2, placeholder="Enter Tamil text here..."),
88
+ outputs=[
89
+ gr.Textbox(label="Translated Text"),
90
+ gr.Textbox(label="Generated Paragraph"),
91
+ gr.Image(type="pil", label="Generated Image")
92
+ ],
93
+ title="Tamil Text Translation, Paragraph Generation, and Image Generation",
94
+ description="Input Tamil text, and this tool will translate it, generate a short paragraph, and create an image based on the translated text."
95
+ )
96
+
97
+ # Launch the Gradio app
98
+ iface.launch()