Added image model

app.py

@@ -2,13 +2,19 @@ import streamlit as st
 from flask.Emotion_spotting_service import _Emotion_spotting_service
 from flask.Genre_spotting_service import _Genre_spotting_service
 from flask.Beat_tracking_service import _Beat_tracking_service
-#from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline
+import tensorflow as tf
 import torch
 import os
 
 emo_list = []
 gen_list = []
 tempo_list = []
+
+physical_devices = tf.config.experimental.list_physical_devices('GPU')
+if len(physical_devices) > 0:
+    tf.config.experimental.set_memory_growth(physical_devices[0], True)
+
 @st.cache_resource
 def load_emo_model():
     emo_service = _Emotion_spotting_service("flask/emotion_model.h5")
@@ -23,11 +29,11 @@ def load_beat_model():
     beat_service = _Beat_tracking_service()
     return beat_service
 
-# @st.cache_resource
-# def load_image_model():
-#      pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",torch_dtype=torch.float16).to("cuda")
-#      pipeline.load_lora_weights("Weights/pytorch_lora_weights.safetensors", weight_name="pytorch_lora_weights.safetensors")
-#      return pipeline
+@st.cache_resource
+def load_image_model():
+     pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",torch_dtype=torch.float16).to("cuda")
+     pipeline.load_lora_weights("Weights/pytorch_lora_weights.safetensors", weight_name="pytorch_lora_weights.safetensors")
+     return pipeline
 
 
 if 'emotion' not in st.session_state:
@@ -42,7 +48,7 @@ if 'beat' not in st.session_state:
 emotion_service = load_emo_model()
 genre_service = load_genre_model()
 beat_service = load_beat_model()
-# image_service = load_image_model()
+image_service = load_image_model()
 
 st.title("Music2Image webpage")
 user_input = st.file_uploader("Upload your wav/mp3 files here", type=["wav","mp3"],key = "file_uploader")
@@ -76,7 +82,7 @@ if st.session_state.emotion != None and st.session_state.genre != None and st.se
         st.caption("Text description of your music file")
         text_output = "A scenic image that describes a " + speed + " pace with a feeling of" + st.session_state.emotion + "."
         st.text(text_output)
-    # if text_output:
-    #    if st.button("Generate image from text description"):
-    #         image = image_service(text_output)
-    #         st.image(image)
+    if text_output:
+       if st.button("Generate image from text description"):
+            image = image_service(text_output)
+            st.image(image)