csaguiar/stable-diffusion-pt

1import os2import torch3import streamlit as st4from diffusers import StableDiffusionPipeline5from transformers import MBart50TokenizerFast, MBartForConditionalGeneration6 7DIFFUSION_MODEL_ID = "runwayml/stable-diffusion-v1-5"8TRANSLATION_MODEL_ID = "Narrativa/mbart-large-50-finetuned-opus-pt-en-translation"  # noqa9DEVICE_NAME = os.getenv("DEVICE_NAME", "cpu")10HUGGING_FACE_TOKEN = os.getenv("HUGGING_FACE_TOKEN")11 12 13def load_translation_models(translation_model_id):14    tokenizer = MBart50TokenizerFast.from_pretrained(15        translation_model_id,16        use_auth_token=HUGGING_FACE_TOKEN17    )18    tokenizer.src_lang = 'pt_XX'19    text_model = MBartForConditionalGeneration.from_pretrained(20        translation_model_id,21        use_auth_token=HUGGING_FACE_TOKEN22    )23 24    return tokenizer, text_model25 26 27def pipeline_generate(diffusion_model_id):28    pipe = StableDiffusionPipeline.from_pretrained(29        diffusion_model_id,30        use_auth_token=HUGGING_FACE_TOKEN31    )32    pipe = pipe.to(DEVICE_NAME)33 34    # Recommended if your computer has < 64 GB of RAM35    pipe.enable_attention_slicing()36 37    return pipe38 39 40def translate(prompt, tokenizer, text_model):41    pt_tokens = tokenizer([prompt], return_tensors="pt")42    en_tokens = text_model.generate(43        **pt_tokens, max_new_tokens=100,44        num_beams=8, early_stopping=True45    )46    en_prompt = tokenizer.batch_decode(en_tokens, skip_special_tokens=True)47 48    return en_prompt[0]49 50 51def generate_image(pipe, prompt):52    # First-time "warmup" pass (see explanation above)53    _ = pipe(prompt, num_inference_steps=1)54 55    return pipe(prompt).images[0]56 57 58def process_prompt(prompt):59    tokenizer, text_model = load_translation_models(TRANSLATION_MODEL_ID)60    prompt = translate(prompt, tokenizer, text_model)61    pipe = pipeline_generate(DIFFUSION_MODEL_ID)62    image = generate_image(pipe, prompt)63    return image64 65 66st.write("# Crie imagens com Stable Diffusion")67prompt_input = st.text_input("Escreva uma descrição da imagem")68 69placeholder = st.empty()70btn = placeholder.button('Processar imagem', disabled=False, key=1)71reload = st.button('Reiniciar', disabled=False)72 73if btn:74    placeholder.button('Processar imagem', disabled=True, key=2)75    image = process_prompt(prompt_input)76    st.image(image)77    placeholder.button('Processar imagem', disabled=False, key=3)78    placeholder.empty()79 80if reload:81    st.experimental_rerun()82