Create app.py

app.py

@@ -0,0 +1,95 @@
+import torch
+from diffusers import StableDiffusionPipeline
+from torch import autocast
+import streamlit as st
+from PIL import Image, ImageEnhance
+import pandas as pd
+import numpy as np
+
+class StableDiffusionLoader:
+    """
+    Stable Diffusion loader and generator class. 
+    Utilises the stable diffusion models from the `Hugging Face`(https://huggingface.co/spaces/stabilityai/stable-diffusion) library
+    Attributes
+    ----------
+    prompt : str
+        a text prompt to use to generate an associated image
+    pretrain_pipe : str
+        a pretrained image diffusion pipeline i.e. CompVis/stable-diffusion-v1-4
+    """
+    def __init__(self, 
+                prompt:str, 
+                pretrain_pipe:str='lfernandopg/mach-5-model-v1'):
+        """
+        Constructs all the necessary attributes for the diffusion class.
+        Parameters
+        ----------
+            prompt : str
+                the prompt to generate the model
+            pretrain_pipe : str
+                the name of the pretrained pipeline
+        """
+        self.prompt = prompt
+        self.pretrain_pipe = pretrain_pipe
+
+        assert isinstance(self.prompt, str), 'Please enter a string into the prompt field'
+        assert isinstance(self.pretrain_pipe, str), 'Please use value such as `CompVis/stable-diffusion-v1-4` for pretrained pipeline'
+
+
+    def generate_image_from_prompt(self, save_location='prompt.jpg', use_token=False,
+                                   verbose=False):
+        """
+        Class method to generate images based on the prompt
+        Parameters
+        ----------
+            save_location : str - defaults to prompt.jpg
+                the location where to save the image generated by the Diffusion Model
+            use_token : bool
+                boolean to see if Hugging Face token should be used
+            verbose : bool
+                boolean that defaults to False, otherwise message printed
+        """
+
+        pipe = StableDiffusionPipeline.from_pretrained(
+            self.pretrain_pipe, 
+            torch_dtype=torch.float16, 
+            )
+        pipe = pipe.to(self.device)
+        with autocast(self.device):
+            image = pipe(self.prompt)[0][0]
+        image.save(save_location)
+        if verbose: 
+            print(f'[INFO] saving image to {save_location}')
+        return image    
+
+    def __str__(self) -> str:
+        return f'[INFO] Generating image for prompt: {self.prompt}'
+
+    def __len__(self):
+        return len(self.prompt)
+
+SAVE_LOCATION = 'prompt.jpg'
+ # Create the page title 
+st.set_page_config(page_title='Diffusion Model generator')
+ # Create page layout
+
+st.title('Image generator using Stable Diffusion')
+st.caption('An app to generate images based on text prompts with a :blue[_Stable Diffusion_] model :sunglasses:')
+
+# Create text prompt
+prompt = st.text_input('Input the prompt desired')
+
+if len(prompt) > 0:
+  st.markdown(f"""
+  This will show an image using **stable diffusion** of the desired {prompt} entered:
+  """)
+  print(prompt)
+  # Create a spinner to show the image is being generated
+  with st.spinner('Generating image based on prompt'):
+    sd = StableDiffusionLoader(prompt)
+    sd.generate_image_from_prompt(save_location=SAVE_LOCATION)
+    st.success('Generated stable diffusion model')
+
+# Open and display the image on the site
+image = Image.open(SAVE_LOCATION)
+st.image(image)