added two models

app.py

@@ -1,38 +1,53 @@
 import gradio as gr
 
-def create_demo(process):
+def generateImage(prompt, n_prompt, modelName):
+
+    return models[modelName].process(prompt, n_prompt)
+
+
+
+
+def create_demo():
+
     with gr.Blocks() as demo:
         with gr.Column():
             prompt = gr.Textbox(label='Prompt')
             n_prompt = gr.Textbox(
                 label='Negative Prompt',
                 value=
-                'low quality, ugly, disfigured, deformed'
+                'ugly, disfigured, deformed'
             )
 
+            modelName = gr.Dropdown(choices = list(models.keys()), 
+                                        label = "Model", 
+                                        value=list(models.keys())[0])
+
             run_button = gr.Button('Run')
-        gr.Markdown("### Visit [Stable Diffusion Art](https://stable-diffusion-art.com/) for tutorials and resources. Read [Model license](https://huggingface.co/spaces/CompVis/stable-diffusion-license) for Stable Diffusion v1.5.")
+        gr.Markdown("### [Stable Diffusion Art](https://stable-diffusion-art.com/) -- tutorials and resources. Read [Model license](https://huggingface.co/spaces/CompVis/stable-diffusion-license).")
         result = gr.Gallery(label='Output',
                             show_label=False,
                             elem_id='gallery').style(columns=1, rows=1, preview=True)        
 
+
         inputs = [
             prompt,
-            n_prompt
+            n_prompt,
+            modelName,
         ]
+
         prompt.submit(
-            fn=process,
+            fn=generateImage,
             inputs=inputs,
             outputs=result
         )
         n_prompt.submit(
-            fn=process,
+            fn=generateImage,
             inputs=inputs,
             outputs=result
         )
 
         run_button.click(
-            fn=process,
+            fn=generateImage,
             inputs=inputs,
             outputs=result
         )
@@ -41,6 +56,10 @@ def create_demo(process):
 
 if __name__ == '__main__':
     from model import Model
-    model = Model()    
-    demo = create_demo(model.process)
+    models = {
+        "Stable Diffusion v1.5": Model("runwayml/stable-diffusion-v1-5"),
+        "Realistic Vision v2.0": Model("SG161222/Realistic_Vision_V2.0"),
+        "Anything v3.0": Model("Linaqruf/anything-v3.0")
+              }
+    demo = create_demo()
     demo.queue().launch()

model.py

@@ -6,6 +6,7 @@ from diffusers import DPMSolverMultistepScheduler
 import torch
 import PIL.Image
 import numpy as np
+import datetime
 
 # Check environment
 print(f"Is CUDA available: {torch.cuda.is_available()}")
@@ -19,11 +20,14 @@ device = "cuda"
 
 
 class Model:
-    def __init__(self):
-        modelID = "runwayml/stable-diffusion-v1-5"
+    def __init__(self, modelID):
+        #modelID = "runwayml/stable-diffusion-v1-5"
+
+        self.modelID = modelID
         self.pipe = StableDiffusionPipeline.from_pretrained(modelID, torch_dtype=torch.float16)
         self.pipe = self.pipe.to(device)
         self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
+        self.pipe.enable_xformers_memory_efficient_attention()
 
         #self.pipe = StableDiffusionPipeline.from_pretrained(modelID)
         #prompt = "a photo of an astronaut riding a horse on mars"
@@ -36,15 +40,20 @@ class Model:
                 num_images:int = 1,
                 num_steps:int = 20,
                 ):
-        seed = np.random.randint(0, np.iinfo(np.int64).max)
+        seed = np.random.randint(0, np.iinfo(np.int32).max)
         generator = torch.Generator(device).manual_seed(seed)
-        return self.pipe(prompt=prompt,
+        now = datetime.datetime.now()
+        print(now)
+        print(self.modelID)
+        with torch.inference_mode():
+            images = self.pipe(prompt=prompt,
                          negative_prompt=negative_prompt,
                          guidance_scale=guidance_scale,
                          num_images_per_prompt=num_images,
                          num_inference_steps=num_steps,
                          generator=generator).images
 
+        return images