Init

.vscode/settings.json

@@ -0,0 +1,16 @@
+{
+    "python.linting.enabled": true,
+    "python.linting.pylintPath": "pylint",
+    "editor.formatOnSave": true,
+    "python.formatting.provider": "yapf", // or "black" here
+    "python.linting.pylintEnabled": true,
+    "python.linting.lintOnSave": true,
+    "python.formatting.yapfArgs": [
+        "--style={based_on_style: google, column_limit: 80, indent_width: 4}"
+      ],
+      "[python]": {
+        "editor.defaultFormatter": "ms-python.black-formatter"
+      },
+    
+      
+}

app.py

@@ -0,0 +1,87 @@
+""" 
+    This file is the main file of the project.
+"""
+
+# imports
+import streamlit as st
+from text_to_image import generate_image
+
+
+def setup():
+    """
+    Streamlit related setup. This has to be run for each page.
+    """
+
+    # hide hamburger menu
+    hide_streamlit_style = """
+    
+    <style>
+    #MainMenu {visibility: hidden;}
+    footer {visibility: hidden;}
+    </style>
+    """
+    st.markdown(hide_streamlit_style, unsafe_allow_html=True)
+
+
+def main():
+    """
+    Main function of the app.
+    """
+
+    setup()
+
+    # title, subheader, and description
+    st.title("Text2Canvas")
+
+    st.subheader("A tool that demonstrates the power of Diffusion")
+
+    # horizontal line and line break
+    st.markdown("<hr>", unsafe_allow_html=True)
+    st.markdown("<br>", unsafe_allow_html=True)
+
+    # sidebar
+    st.sidebar.title("Navigation")
+    mode = st.sidebar.radio("Select a mode", ["Home", "Text2Image", "Feature2Sprite"])
+    # st.sidebar.write("Select a mode to get started")
+
+    # main content
+
+    if mode == "Home":
+        st.write(
+            """
+        This tool is a demonstration of the power of Diffusion. It helps you generate images from text. There are two modes:
+        1. **Text2Image**: This mode generates high quality image based on a given prompt. It uses a pretrained model.
+        2. **Feature2Sprite**: This mode generates 16*16 images of sprites based on a combination of features. It uses a custom model trained on a dataset of sprites.
+        
+        To get started, select a mode from the sidebar.
+        """
+        )
+
+    elif mode == "Text2Image":
+        st.write(
+            """
+        This mode generates high quality image based on a given prompt. It uses a pretrained model from huggingface.
+        """
+        )
+
+        form = st.form(key="my_form")
+
+        prompt = form.text_input("Enter a prompt", value="A painting of a cat")
+
+        submit_button = form.form_submit_button(label="Generate")
+
+        if submit_button:
+            st.write("Generating image...")
+            image = generate_image(prompt)
+            st.image(image, caption="Generated Image", use_column_width=True)
+
+    elif mode == "Feature2Sprite":
+        st.write(
+            """
+        This mode generates 16*16 images of sprites based on a combination of features. It uses a custom model trained on a dataset of sprites.
+        """
+        )
+
+
+if __name__ == "__main__":
+    main()

constants.py

@@ -0,0 +1,5 @@
+""" 
+    This file contains all the constants used in the project.
+"""
+
+MODEL_ID = "stabilityai/stable-diffusion-2-1"

requirements.txt

@@ -0,0 +1,4 @@
+diffusers==0.19.3
+streamlit==1.25.0
+torch==2.0.1
+transformers==4.31.0

text_to_image.py

@@ -0,0 +1,33 @@
+""" 
+    This file contains the code for the text to image generation using the Stable Diffusion model.
+"""
+
+import torch
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+
+from constants import MODEL_ID
+
+
+def generate_image(prompt: str) -> torch.Tensor:
+    """
+    Generates an image from a given prompt.
+
+    Args:
+        prompt (str): The prompt to generate the image from.
+
+    Returns:
+        torch.Tensor: The generated image.
+    """
+    # load model
+
+    pipe = StableDiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.float16)
+    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+    if torch.cuda.is_available():
+        pipe = pipe.to("cuda")  # move model to GPU if available
+        print("Model loaded successfully and moved to GPU.")
+    else:
+        print("Model loaded successfully on CPU.")
+
+    image = pipe(prompt).images[0]
+
+    return image