iohanngrig/text2image

1import streamlit as st2from io import BytesIO3from typing import Literal4from diffusers import StableDiffusionPipeline5import torch6import time7 8seed = 429generator = torch.manual_seed(seed)10 11NUM_ITERS_TO_RUN = 212NUM_INFERENCE_STEPS = 2013NUM_IMAGES_PER_PROMPT = 114 15 16def text2image(17    prompt: str,18    repo_id: Literal[19        "dreamlike-art/dreamlike-photoreal-2.0",20        "hakurei/waifu-diffusion",21        "prompthero/openjourney",22        "stabilityai/stable-diffusion-2-1",23        "runwayml/stable-diffusion-v1-5",24        "nota-ai/bk-sdm-small",25        "CompVis/stable-diffusion-v1-4",26    ],27):28    start = time.time()29    if torch.cuda.is_available():30        print("Using GPU")31        pipeline = StableDiffusionPipeline.from_pretrained(32            repo_id,33            torch_dtype=torch.float16,34            use_safetensors=True,35        ).to("cuda")36    else:37        print("Using CPU")38        pipeline = StableDiffusionPipeline.from_pretrained(39            repo_id,40            torch_dtype=torch.float32,41            use_safetensors=True,42        )43 44    for _ in range(NUM_ITERS_TO_RUN):45        images = pipeline(46            prompt,47            num_inference_steps=NUM_INFERENCE_STEPS,48            generator=generator,49            num_images_per_prompt=NUM_IMAGES_PER_PROMPT,50        ).images51    end = time.time()52    return images[0], start, end53 54def app():55    st.header("Text-to-image Web App")56    st.subheader("Powered by Hugging Face")57    user_input = st.text_area(58        "Enter your text prompt below and click the button to submit."59    )60 61    option = st.selectbox(62        "Select model (in order of processing time)",63        (64            "nota-ai/bk-sdm-small",65            "CompVis/stable-diffusion-v1-4",66            "runwayml/stable-diffusion-v1-5",67            "prompthero/openjourney",68            "hakurei/waifu-diffusion",69            "stabilityai/stable-diffusion-2-1",70            "dreamlike-art/dreamlike-photoreal-2.0",71        ),72    )73 74    with st.form("my_form"):75        submit = st.form_submit_button(label="Submit text prompt")76 77    if submit:78        with st.spinner(text="Generating image ... It may take up to 20 minutes."):79            im, start, end = text2image(prompt=user_input, repo_id=option)80            buf = BytesIO()81            im.save(buf, format="PNG")82            byte_im = buf.getvalue()83            hours, rem = divmod(end - start, 3600)84            minutes, seconds = divmod(rem, 60)85            st.success(86                "Processing time: {:0>2}:{:0>2}:{:05.2f}.".format(87                    int(hours), int(minutes), seconds88                )89            )90            st.image(im)91            st.download_button(92                label="Click here to download",93                data=byte_im,94                file_name="generated_image.png",95                mime="image/png",96            )97 98if __name__ == "__main__":99    app()