coztomate commited on
Commit
6befe96
1 Parent(s): 73cc78e

changed model to preloaded in beginning

Browse files
Files changed (1) hide show
  1. app.py +42 -43
app.py CHANGED
@@ -10,7 +10,6 @@ from openai import OpenAI
10
  import openai
11
  from diffusers import StableDiffusionPipeline
12
 
13
-
14
  # Initialize session state variables
15
  if 'simplified_text' not in st.session_state:
16
  st.session_state['simplified_text'] = ''
@@ -33,7 +32,9 @@ if 'image_from_simplified_text' not in st.session_state:
33
  if 'image_from_press_text' not in st.session_state:
34
  st.session_state['image_from_press_text'] = None
35
 
 
36
 
 
37
  # Define model and tokenizer names for the text simplification model
38
  model_name = "mrm8488/t5-small-finetuned-text-simplification"
39
  tokenizer_name = "mrm8488/t5-small-finetuned-text-simplification"
@@ -47,7 +48,8 @@ if 'model' not in st.session_state or 'tokenizer' not in st.session_state:
47
  # Use the model from session state
48
  simplifier = st.session_state['simplifier']
49
 
50
- # Function to load CLIP model
 
51
  def load_clip_model():
52
  model_clip, _, transform_clip = open_clip.create_model_and_transforms(
53
  model_name="coca_ViT-L-14",
@@ -55,31 +57,45 @@ def load_clip_model():
55
  )
56
  return model_clip, transform_clip
57
 
58
- # Function to generate a caption for the uploaded image
 
 
 
59
  def generate_caption(image_path):
60
- # Load the CLIP model if it hasn't been loaded yet
61
- if st.session_state['model_clip'] is None or st.session_state['transform_clip'] is None:
62
- st.session_state['model_clip'], st.session_state['transform_clip'] = load_clip_model()
63
-
64
- # Load and preprocess the uploaded image
65
  im = Image.open(image_path).convert("RGB")
66
- im = st.session_state['transform_clip'](im).unsqueeze(0)
67
 
68
  # Generate a caption for the image
69
  with torch.no_grad(), torch.cuda.amp.autocast():
70
- generated = st.session_state['model_clip'].generate(im)
71
 
72
  new_caption = open_clip.decode(generated[0]).split("<end_of_text>")[0].replace("<start_of_text>", "")[:-2]
73
  return new_caption
74
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
75
 
 
76
 
77
  # Create a Streamlit app
78
- st.title("ARTSPEAK s i m p l i f i e r")
79
 
80
  # Display an image from the local file system
81
  logo_path = 'logo_artspeak.png' # Replace with your image path
82
- st.image(logo_path, caption='Logo', use_column_width='auto')
83
 
84
  st.markdown("---")
85
 
@@ -118,14 +134,14 @@ if st.button("Simplify"):
118
 
119
  # Display the simplified text from session state
120
  if st.session_state['simplified_text']:
121
- st.write("Simplified Text:")
122
  st.write(st.session_state['simplified_text'])
123
 
124
  st.markdown("---")
125
 
126
  ####Get new caption
127
  # Modify the 'Get Caption' button section
128
- if st.button("Get Caption"):
129
  if uploaded_image is not None:
130
  # Generate the caption
131
  caption = generate_caption(uploaded_image)
@@ -136,7 +152,7 @@ if st.button("Get Caption"):
136
 
137
  # Display the new caption from session state
138
  if st.session_state['new_caption']:
139
- st.write("New Caption for Artwork:")
140
  st.write(st.session_state['new_caption'])
141
 
142
  st.markdown("---")
@@ -178,7 +194,7 @@ if st.button("Generate Press Text from New Caption"):
178
 
179
  # Display the generated press text from new caption
180
  if st.session_state['message_content_from_caption']:
181
- st.write("Generated Press Text from New Caption:")
182
  st.write(st.session_state['message_content_from_caption'])
183
 
184
  # Button to generate press text from simplified text
@@ -201,54 +217,37 @@ st.markdown("---")
201
  ############
202
  ##Diffusor##
203
  ############
204
- # Load Stable Diffusion model
205
- def load_diffusion_model():
206
- pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
207
- pipe = pipe.to("cuda")
208
- return pipe
209
 
210
- # Function to generate an image and show a notification while processing
211
- def generate_image(pipe, prompt, notification_placeholder):
212
- with notification_placeholder.container():
213
- st.text('Generating image...')
214
- image = pipe(prompt).images[0]
215
- st.empty() # Clear the notification after the process is done
216
- return image
217
 
218
- # Button and notification for generating image from new caption
219
- if st.button("Generate Image from New Caption"):
220
- notification_placeholder = st.empty()
221
  if st.session_state['new_caption']:
222
- pipe = load_diffusion_model()
223
  prompt_caption = f"contemporary art of {st.session_state['new_caption']}"
224
- st.session_state['image_from_caption'] = generate_image(pipe, prompt_caption, notification_placeholder)
225
 
226
  # Display the image generated from new caption
227
  if st.session_state['image_from_caption'] is not None:
228
  st.image(st.session_state['image_from_caption'], caption="Image from New Caption", use_column_width=True)
229
 
230
-
231
  # Button to generate image from simplified text
232
  if st.button("Generate Image from Simplified Text"):
233
- notification_placeholder = st.empty()
234
  if st.session_state['simplified_text']:
235
- pipe = load_diffusion_model()
236
  prompt_summary = f"contemporary art of {st.session_state['simplified_text']}"
237
- st.session_state['image_from_simplified_text'] = generate_image(pipe, prompt_summary, notification_placeholder)
238
 
239
  # Display the image generated from simplified text
240
  if st.session_state['image_from_simplified_text'] is not None:
241
  st.image(st.session_state['image_from_simplified_text'], caption="Image from Simplified Text", use_column_width=True)
242
 
243
-
244
  # Button to generate image from press text
245
- if st.button("Generate Image from Press Text"):
246
- notification_placeholder = st.empty()
247
  if st.session_state['message_content_from_simplified_text']:
248
- pipe = load_diffusion_model()
249
  prompt_press_text = f"contemporary art of {st.session_state['message_content_from_simplified_text']}"
250
- st.session_state['image_from_press_text'] = generate_image(pipe, prompt_press_text, notification_placeholder)
251
 
252
  # Display the image generated from press text
253
  if st.session_state['image_from_press_text'] is not None:
254
- st.image(st.session_state['image_from_press_text'], caption="Image from Press Text", use_column_width=True)
 
 
 
 
10
  import openai
11
  from diffusers import StableDiffusionPipeline
12
 
 
13
  # Initialize session state variables
14
  if 'simplified_text' not in st.session_state:
15
  st.session_state['simplified_text'] = ''
 
32
  if 'image_from_press_text' not in st.session_state:
33
  st.session_state['image_from_press_text'] = None
34
 
35
+ ######loading models########
36
 
37
+ ####loading simplifier model#####
38
  # Define model and tokenizer names for the text simplification model
39
  model_name = "mrm8488/t5-small-finetuned-text-simplification"
40
  tokenizer_name = "mrm8488/t5-small-finetuned-text-simplification"
 
48
  # Use the model from session state
49
  simplifier = st.session_state['simplifier']
50
 
51
+ ####loading clip model#####
52
+ # Function to load the CLIP model
53
  def load_clip_model():
54
  model_clip, _, transform_clip = open_clip.create_model_and_transforms(
55
  model_name="coca_ViT-L-14",
 
57
  )
58
  return model_clip, transform_clip
59
 
60
+ if 'loaded_clip_model' not in st.session_state or 'loaded_transform_clip' not in st.session_state:
61
+ st.session_state['loaded_clip_model'], st.session_state['loaded_transform_clip'] = load_clip_model()
62
+
63
+ # Function to generate a caption using the preloaded CLIP model
64
  def generate_caption(image_path):
 
 
 
 
 
65
  im = Image.open(image_path).convert("RGB")
66
+ im = st.session_state['loaded_transform_clip'](im).unsqueeze(0)
67
 
68
  # Generate a caption for the image
69
  with torch.no_grad(), torch.cuda.amp.autocast():
70
+ generated = st.session_state['loaded_clip_model'].generate(im)
71
 
72
  new_caption = open_clip.decode(generated[0]).split("<end_of_text>")[0].replace("<start_of_text>", "")[:-2]
73
  return new_caption
74
 
75
+ ###loading diffusion model
76
+ # Function to load the Stable Diffusion model
77
+ def load_diffusion_model():
78
+ pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
79
+ pipe = pipe.to("cuda")
80
+ return pipe
81
+
82
+ # Initialize the model at the start and store it in the session state
83
+ if 'loaded_model' not in st.session_state:
84
+ st.session_state['loaded_model'] = load_diffusion_model()
85
+
86
+ # Function to generate an image using the preloaded model
87
+ def generate_image(prompt):
88
+ image = st.session_state['loaded_model'](prompt).images[0]
89
+ return image
90
 
91
+ ################################################
92
 
93
  # Create a Streamlit app
94
+ st.title("ARTSPEAK > s i m p l i f i e r")
95
 
96
  # Display an image from the local file system
97
  logo_path = 'logo_artspeak.png' # Replace with your image path
98
+ st.image(logo_path, use_column_width='auto')
99
 
100
  st.markdown("---")
101
 
 
134
 
135
  # Display the simplified text from session state
136
  if st.session_state['simplified_text']:
137
+ st.write("Simplified Original Text:")
138
  st.write(st.session_state['simplified_text'])
139
 
140
  st.markdown("---")
141
 
142
  ####Get new caption
143
  # Modify the 'Get Caption' button section
144
+ if st.button("Get New Caption"):
145
  if uploaded_image is not None:
146
  # Generate the caption
147
  caption = generate_caption(uploaded_image)
 
152
 
153
  # Display the new caption from session state
154
  if st.session_state['new_caption']:
155
+ st.write("New Caption for this Artwork:")
156
  st.write(st.session_state['new_caption'])
157
 
158
  st.markdown("---")
 
194
 
195
  # Display the generated press text from new caption
196
  if st.session_state['message_content_from_caption']:
197
+ st.write("Generated Press Text from New Caption of Artwork:")
198
  st.write(st.session_state['message_content_from_caption'])
199
 
200
  # Button to generate press text from simplified text
 
217
  ############
218
  ##Diffusor##
219
  ############
 
 
 
 
 
220
 
 
 
 
 
 
 
 
221
 
222
+ # Example button to generate image from new caption
223
+ if st.button("Generate Image from New Caption of Artwork"):
 
224
  if st.session_state['new_caption']:
 
225
  prompt_caption = f"contemporary art of {st.session_state['new_caption']}"
226
+ st.session_state['image_from_caption'] = generate_image(prompt_caption)
227
 
228
  # Display the image generated from new caption
229
  if st.session_state['image_from_caption'] is not None:
230
  st.image(st.session_state['image_from_caption'], caption="Image from New Caption", use_column_width=True)
231
 
 
232
  # Button to generate image from simplified text
233
  if st.button("Generate Image from Simplified Text"):
 
234
  if st.session_state['simplified_text']:
 
235
  prompt_summary = f"contemporary art of {st.session_state['simplified_text']}"
236
+ st.session_state['image_from_simplified_text'] = generate_image(pipe, prompt_summary)
237
 
238
  # Display the image generated from simplified text
239
  if st.session_state['image_from_simplified_text'] is not None:
240
  st.image(st.session_state['image_from_simplified_text'], caption="Image from Simplified Text", use_column_width=True)
241
 
 
242
  # Button to generate image from press text
243
+ if st.button("Generate Image from new Press Text"):
 
244
  if st.session_state['message_content_from_simplified_text']:
 
245
  prompt_press_text = f"contemporary art of {st.session_state['message_content_from_simplified_text']}"
246
+ st.session_state['image_from_press_text'] = generate_image(pipe, prompt_press_text)
247
 
248
  # Display the image generated from press text
249
  if st.session_state['image_from_press_text'] is not None:
250
+ st.image(st.session_state['image_from_press_text'], caption="Image from Press Text", use_column_width=True)
251
+
252
+
253
+ st.markdown("---")