added diffusion model

app.py

@@ -1,5 +1,8 @@
 # Import convention
 import streamlit as st
+from diffusers import DiffusionPipeline
+import matplotlib.pyplot as plt
+import torch
 
 organ = st.selectbox('Organ', ['Brain', 'Thorax'], index=None)
 modality = st.selectbox('Modality', ['Magnetic Resonance Imaging', 'Computed Tomography'], index=None)
@@ -8,4 +11,11 @@ style = st.selectbox('Style', ['Picasso', 'Van Gogh'], index=None)
 prompt_lst = [organ, modality, style]
 if None not in prompt_lst:
     prompt = ','.join(prompt_lst) 
-    print(prompt)
+    print(prompt)
+
+    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16")
+    pipe.to("cuda")
+
+    prompt += " high resolution, photorealistic"
+    image = pipe(prompt=prompt).images[0]
+    st.image(image)