Update app.py

app.py

@@ -12,6 +12,9 @@ image_to_text = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-capti
 # Load the text mask pipeline
 generate_mask = pipeline("fill-mask", model="google-bert/bert-base-uncased")
 
+# Load the text generation pipeline
+extend_text = pipeline("text-generation", model="pranavpsv/genre-story-generator-v2")
+
 # Load the text-to-image model
 #text_to_image = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
 
@@ -66,7 +69,8 @@ def generate_mask_from_result(text):
     else:
         result = output[0]['token_str'] if output else None
 
-    return result
+    extended_text = extend_text(f"A child with {text} and {result} ")
+    return extended_text
 
 
 if __name__ == "__main__":