Update app.py

app.py

@@ -286,7 +286,23 @@ def captioning(img):
     )
 
     generated_text = processor5.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+
+    generated_idsb = model5.generate(
+        **inputsa,
+        do_sample=True,
+        num_beams=1,
+        max_length=16,
+        min_length=12,
+        top_p=0.9,
+        repetition_penalty=1.0,
+        length_penalty=2.0,
+        temperature=0.5,
+    )
+
+    generated_textb = processor5.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+
     generated_text = generated_text.replace(cap_prompt, "").strip() #Or could try .split(prompt, 1)[-1].strip()
+    generated_textb = generated_textb.replace(cap_prompt, "").strip() #Or could try .split(prompt, 1)[-1].strip()
     output_prompt.append(generated_text)
     print(generated_text)
     # Loop through prompts array:
@@ -331,7 +347,7 @@ def captioning(img):
     output_prompt.append(response_text)
     '''
     output_prompt = " ".join(output_prompt)
-    return output_prompt
+    return output_prompt, generated_textb
 
 def flatten_and_stringify(data):
     return [str(item) for sublist in data if isinstance(sublist, list) for item in flatten_and_stringify(sublist) ] + [str(item) for item in data if not isinstance(item, list)]
@@ -433,19 +449,22 @@ def generate_30(
     image_paths=[]
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
+    capt=[]
     if latent_file is not None:  # Check if a latent file is provided
         sd_image_a = Image.open(latent_file.name).convert('RGB')
         sd_image_a.resize((768,768), Image.LANCZOS)
         #sd_image_a.resize((height,width), Image.LANCZOS)
         caption=[]
         #caption.append(captioning(sd_image_a))
-        prompt1=captioning(sd_image_a)
+        prompt1, cap=captioning(sd_image_a)
+        capt.append(cap)
         if latent_file_2 is not None:  # Check if a latent file is provided
             sd_image_b = Image.open(latent_file_2.name).convert('RGB')
             #sd_image_b.resize((height,width), Image.LANCZOS)
             sd_image_b.resize((768,768), Image.LANCZOS)
             #caption.append(captioning(sd_image_b))
-            prompt2=captioning(sd_image_b)
+            prompt2, cap=captioning(sd_image_b)
+            capt.append(cap)
         else:
             sd_image_b = None
         if latent_file_3 is not None:  # Check if a latent file is provided
@@ -453,7 +472,8 @@ def generate_30(
             #sd_image_c.resize((height,width), Image.LANCZOS)
             sd_image_c.resize((768,768), Image.LANCZOS)
             #caption.append(captioning(sd_image_c))
-            prompt3=captioning(sd_image_c)
+            prompt3, cap=captioning(sd_image_c)
+            capt.append(cap)
         else:
             sd_image_c = None
         if latent_file_4 is not None:  # Check if a latent file is provided
@@ -461,7 +481,8 @@ def generate_30(
             #sd_image_d.resize((height,width), Image.LANCZOS)
             sd_image_d.resize((768,768), Image.LANCZOS)
             #caption.append(captioning(sd_image_d))
-            prompt4=captioning(sd_image_d)
+            prompt4, cap=captioning(sd_image_d)
+            capt.append(cap)
         else:
             sd_image_d = None
         if latent_file_5 is not None:  # Check if a latent file is provided
@@ -469,7 +490,8 @@ def generate_30(
             #sd_image_e.resize((height,width), Image.LANCZOS)
             sd_image_e.resize((768,768), Image.LANCZOS)
             #caption.append(captioning(sd_image_e))
-            prompt5=captioning(sd_image_e)
+            prompt5, cap=captioning(sd_image_e)
+            capt.append(cap)
         else:
             sd_image_e = None
         timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -481,6 +503,7 @@ def generate_30(
         print("-- CURRENT PROMPT AFTER .join --")
         print(prompt)
         captions = " ".join(caption)
+        capt = " ".join(capt)
         print(captions)
         print("-- not generating further caption --")
         global model5
@@ -494,7 +517,7 @@ def generate_30(
         new_prompt = prompt + ' ' + captions
         print("-- ------------ --")
         print("-- FINAL PROMPT --")
-        print(prompt)
+        print(capt)
         print("-- FINAL PROMPT --")
         print("-- ------------ --")
         gc.collect()
@@ -577,19 +600,22 @@ def generate_60(
     image_paths=[]
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
+    capt=[]
     if latent_file is not None:  # Check if a latent file is provided
         sd_image_a = Image.open(latent_file.name).convert('RGB')
         sd_image_a.resize((768,768), Image.LANCZOS)
         #sd_image_a.resize((height,width), Image.LANCZOS)
         caption=[]
         #caption.append(captioning(sd_image_a))
-        prompt1=captioning(sd_image_a)
+        prompt1, cap=captioning(sd_image_a)
+        capt.append(cap)
         if latent_file_2 is not None:  # Check if a latent file is provided
             sd_image_b = Image.open(latent_file_2.name).convert('RGB')
             #sd_image_b.resize((height,width), Image.LANCZOS)
             sd_image_b.resize((768,768), Image.LANCZOS)
             #caption.append(captioning(sd_image_b))
-            prompt2=captioning(sd_image_b)
+            prompt2, cap=captioning(sd_image_b)
+            capt.append(cap)
         else:
             sd_image_b = None
         if latent_file_3 is not None:  # Check if a latent file is provided
@@ -597,7 +623,8 @@ def generate_60(
             #sd_image_c.resize((height,width), Image.LANCZOS)
             sd_image_c.resize((768,768), Image.LANCZOS)
             #caption.append(captioning(sd_image_c))
-            prompt3=captioning(sd_image_c)
+            prompt3, cap=captioning(sd_image_c)
+            capt.append(cap)
         else:
             sd_image_c = None
         if latent_file_4 is not None:  # Check if a latent file is provided
@@ -605,7 +632,8 @@ def generate_60(
             #sd_image_d.resize((height,width), Image.LANCZOS)
             sd_image_d.resize((768,768), Image.LANCZOS)
             #caption.append(captioning(sd_image_d))
-            prompt4=captioning(sd_image_d)
+            prompt4, cap=captioning(sd_image_d)
+            capt.append(cap)
         else:
             sd_image_d = None
         if latent_file_5 is not None:  # Check if a latent file is provided
@@ -613,7 +641,8 @@ def generate_60(
             #sd_image_e.resize((height,width), Image.LANCZOS)
             sd_image_e.resize((768,768), Image.LANCZOS)
             #caption.append(captioning(sd_image_e))
-            prompt5=captioning(sd_image_e)
+            prompt5, cap=captioning(sd_image_e)
+            capt.append(cap)
         else:
             sd_image_e = None
         timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -625,6 +654,7 @@ def generate_60(
         print("-- CURRENT PROMPT AFTER .join --")
         print(prompt)
         captions = " ".join(caption)
+        capt = " ".join(capt)
         print(captions)
         print("-- not generating further caption --")
         global model5
@@ -638,7 +668,7 @@ def generate_60(
         new_prompt = prompt + ' ' + captions
         print("-- ------------ --")
         print("-- FINAL PROMPT --")
-        print(prompt)
+        print(capt)
         print("-- FINAL PROMPT --")
         print("-- ------------ --")
         gc.collect()
@@ -721,19 +751,22 @@ def generate_90(
     image_paths=[]
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
+    capt=[]
     if latent_file is not None:  # Check if a latent file is provided
         sd_image_a = Image.open(latent_file.name).convert('RGB')
         sd_image_a.resize((768,768), Image.LANCZOS)
         #sd_image_a.resize((height,width), Image.LANCZOS)
         caption=[]
         #caption.append(captioning(sd_image_a))
-        prompt1=captioning(sd_image_a)
+        prompt1, cap=captioning(sd_image_a)
+        capt.append(cap)
         if latent_file_2 is not None:  # Check if a latent file is provided
             sd_image_b = Image.open(latent_file_2.name).convert('RGB')
             #sd_image_b.resize((height,width), Image.LANCZOS)
             sd_image_b.resize((768,768), Image.LANCZOS)
             #caption.append(captioning(sd_image_b))
-            prompt2=captioning(sd_image_b)
+            prompt2, cap=captioning(sd_image_b)
+            capt.append(cap)
         else:
             sd_image_b = None
         if latent_file_3 is not None:  # Check if a latent file is provided
@@ -741,7 +774,8 @@ def generate_90(
             #sd_image_c.resize((height,width), Image.LANCZOS)
             sd_image_c.resize((768,768), Image.LANCZOS)
             #caption.append(captioning(sd_image_c))
-            prompt3=captioning(sd_image_c)
+            prompt3, cap=captioning(sd_image_c)
+            capt.append(cap)
         else:
             sd_image_c = None
         if latent_file_4 is not None:  # Check if a latent file is provided
@@ -749,7 +783,8 @@ def generate_90(
             #sd_image_d.resize((height,width), Image.LANCZOS)
             sd_image_d.resize((768,768), Image.LANCZOS)
             #caption.append(captioning(sd_image_d))
-            prompt4=captioning(sd_image_d)
+            prompt4, cap=captioning(sd_image_d)
+            capt.append(cap)
         else:
             sd_image_d = None
         if latent_file_5 is not None:  # Check if a latent file is provided
@@ -757,7 +792,8 @@ def generate_90(
             #sd_image_e.resize((height,width), Image.LANCZOS)
             sd_image_e.resize((768,768), Image.LANCZOS)
             #caption.append(captioning(sd_image_e))
-            prompt5=captioning(sd_image_e)
+            prompt5, cap=captioning(sd_image_e)
+            capt.append(cap)
         else:
             sd_image_e = None
         timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -769,6 +805,7 @@ def generate_90(
         print("-- CURRENT PROMPT AFTER .join --")
         print(prompt)
         captions = " ".join(caption)
+        capt = " ".join(capt)
         print(captions)
         print("-- not generating further caption --")
         global model5
@@ -782,7 +819,7 @@ def generate_90(
         new_prompt = prompt + ' ' + captions
         print("-- ------------ --")
         print("-- FINAL PROMPT --")
-        print(prompt)
+        print(capt)
         print("-- FINAL PROMPT --")
         print("-- ------------ --")
         gc.collect()