File size: 2,289 Bytes
1437058
2d8ff37
d550533
eea6ac5
d16c1e4
9fd10ba
d7164de
2d8ff37
 
 
 
9bdb949
2d8ff37
9bdb949
9fd10ba
2d8ff37
d7164de
2d8ff37
 
 
 
9bdb949
2d8ff37
5ca4bcd
2d8ff37
 
 
 
 
9607ff2
c1732d5
2d8ff37
 
 
 
4d86cc2
2d8ff37
 
 
1437058
2d8ff37
 
 
 
1437058
2d8ff37
 
 
 
1437058
2d8ff37
 
 
 
9607ff2
 
d16c1e4
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
import streamlit as st
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
from diffusers import StableDiffusionPipeline
import torch

@st.cache_resource
def load_all_models():
    # Load IndicTrans2 Tamil-to-English model
    trans_model_id = "ai4bharat/indictrans2-indic-en-dist-200M"
    trans_tokenizer = AutoTokenizer.from_pretrained(trans_model_id, trust_remote_code=True)
    trans_model = AutoModelForSeq2SeqLM.from_pretrained(trans_model_id, trust_remote_code=True)

    # Load English text generation model (you can use GPT2 or Falcon, etc.)
    text_gen = pipeline("text-generation", model="gpt2")

    # Load Stable Diffusion for image generation
    img_pipe = StableDiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        torch_dtype=torch.float16,
        revision="fp16"
    ).to("cuda" if torch.cuda.is_available() else "cpu")

    return trans_tokenizer, trans_model, text_gen, img_pipe

def translate_text(text, tokenizer, model):
    input_text = f"translate Tamil to English: {text}"
    inputs = tokenizer(input_text, return_tensors="pt", padding=True)
    outputs = model.generate(**inputs, max_length=128)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def main():
    st.title("Multimodal Tamil to Image Generator 🚀")
    st.markdown("Enter Tamil text, we translate it to English, continue the sentence, and generate an image!")

    user_input = st.text_area("Enter Tamil text:", "")

    if st.button("Generate"):
        with st.spinner("Loading models..."):
            tokenizer, model, text_gen, img_pipe = load_all_models()

        with st.spinner("Translating to English..."):
            english_text = translate_text(user_input, tokenizer, model)
            st.subheader("Translated English:")
            st.write(english_text)

        with st.spinner("Generating continuation..."):
            continuation = text_gen(english_text, max_length=50, do_sample=True)[0]['generated_text']
            st.subheader("Generated Text:")
            st.write(continuation)

        with st.spinner("Generating Image..."):
            image = img_pipe(continuation).images[0]
            st.subheader("Generated Image:")
            st.image(image)

if __name__ == "__main__":
    main()