first app with image generation in portuguese

app.py

@@ -1,4 +1,77 @@
 import streamlit as st
+from diffusers import StableDiffusionPipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+DIFFUSION_MODEL_ID = "runwayml/stable-diffusion-v1-5"
+TRANSLATION_MODEL_ID = "Narrativa/mbart-large-50-finetuned-opus-pt-en-translation"  # noqa
+
+
+def load_translation_models(translation_model_id):
+    tokenizer = AutoTokenizer.from_pretrained(
+        translation_model_id,
+        use_auth_token=True
+    )
+    text_model = AutoModelForSeq2SeqLM.from_pretrained(
+        translation_model_id,
+        use_auth_token=True
+    )
+
+    return tokenizer, text_model
+
+
+def pipeline_generate(diffusion_model_id):
+    pipe = StableDiffusionPipeline.from_pretrained(
+        diffusion_model_id,
+        use_auth_token=True
+    )
+    pipe = pipe.to("mps")
+
+    # Recommended if your computer has < 64 GB of RAM
+    pipe.enable_attention_slicing()
+
+    return pipe
+
+
+def translate(prompt, tokenizer, text_model):
+    pt_tokens = tokenizer([prompt], return_tensors="pt")
+    en_tokens = text_model.generate(
+        **pt_tokens, max_new_tokens=100,
+        num_beams=8, early_stopping=True
+    )
+    en_prompt = tokenizer.batch_decode(en_tokens, skip_special_tokens=True)
+    print(f"translation: [PT] {prompt} -> [EN] {en_prompt[0]}")
+
+    return en_prompt[0]
+
+
+def generate_image(pipe, prompt):
+    # First-time "warmup" pass (see explanation above)
+    _ = pipe(prompt, num_inference_steps=1)
+
+    return pipe(prompt).images[0]
+
+
+def process_prompt(prompt):
+    tokenizer, text_model = load_translation_models(TRANSLATION_MODEL_ID)
+    prompt = translate(prompt, tokenizer, text_model)
+    pipe = pipeline_generate(DIFFUSION_MODEL_ID)
+    image = generate_image(pipe, prompt)
+    return image
+
+
+st.write("# Crie imagens com Stable Diffusion")
+prompt_input = st.text_input("Escreva uma descrição da imagem")
+
+placeholder = st.empty()
+btn = placeholder.button('Processar imagem', disabled=False, key=1)
+reload = st.button('Reiniciar', disabled=False)
+
+if btn:
+    placeholder.button('Processar imagem', disabled=True, key=2)
+    image = process_prompt(prompt_input)
+    st.image(image)
+    placeholder.button('Processar imagem', disabled=False, key=3)
+    placeholder.empty()
+
+if reload:
+    st.experimental_rerun()