Update ImageGenerator.py

ImageGenerator.py

@@ -124,13 +124,25 @@ pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float
 def generate_image(prompt):
     # Get the part after the dash
     print('PROMPT TO INTRODUCE TO STABLDIFFUSION:', prompt)
-    if "-" in prompt:
-        prompt_to_send = prompt.split("-", 1)[-1].strip()
-    else:
-        prompt_to_send = prompt.strip()
+    prompt_to_send = extract_assistant_response(prompt)
+    print('INTRODUCED: ', prompt_to_send)
+    # if "-" in prompt:
+    #     prompt_to_send = prompt.split("-", 1)[-1].strip()
+    # else:
+    #     prompt_to_send = prompt.strip()
     # prompt = "a photo of an astronaut riding a horse on mars"
     image = pipe(prompt_to_send).images[0]  
     return image
+
+def extract_assistant_response(decoded_text):
+    if "<|im_start|>assistant" in decoded_text:
+        part = decoded_text.split("<|im_start|>assistant", 1)[1]
+        part = part.split("<|im_end|>", 1)[0]
+        return part.strip()
+    return decoded_text.strip()
+
+
+    
 # image.save("astronaut_rides_horse.png")