import os import torch import streamlit as st from diffusers import StableDiffusionPipeline from transformers import MBart50TokenizerFast, MBartForConditionalGeneration DIFFUSION_MODEL_ID = "runwayml/stable-diffusion-v1-5" TRANSLATION_MODEL_ID = "Narrativa/mbart-large-50-finetuned-opus-pt-en-translation" # noqa DEVICE_NAME = os.getenv("DEVICE_NAME", "cpu") HUGGING_FACE_TOKEN = os.getenv("HUGGING_FACE_TOKEN") def load_translation_models(translation_model_id): tokenizer = MBart50TokenizerFast.from_pretrained( translation_model_id, use_auth_token=HUGGING_FACE_TOKEN ) tokenizer.src_lang = 'pt_XX' text_model = MBartForConditionalGeneration.from_pretrained( translation_model_id, use_auth_token=HUGGING_FACE_TOKEN ) return tokenizer, text_model def pipeline_generate(diffusion_model_id): pipe = StableDiffusionPipeline.from_pretrained( diffusion_model_id, torch_dtype=torch.float16, revision="fp16", use_auth_token=HUGGING_FACE_TOKEN ) pipe = pipe.to(DEVICE_NAME) # 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) 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()