Spaces:
Sleeping
Sleeping
created app.py file
Browse files
app.py
ADDED
@@ -0,0 +1,100 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from fastapi import FastAPI, HTTPException
|
2 |
+
from pydantic import BaseModel
|
3 |
+
from langchain_huggingface.llms import HuggingFacePipeline
|
4 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
|
5 |
+
from langchain_core.callbacks import CallbackManager, StreamingStdOutCallbackHandler
|
6 |
+
from langchain_core.prompts import ChatPromptTemplate
|
7 |
+
|
8 |
+
import os
|
9 |
+
import fal_client
|
10 |
+
|
11 |
+
# FastAPI app
|
12 |
+
app = FastAPI()
|
13 |
+
|
14 |
+
# Set the environment variable
|
15 |
+
os.environ['FAL_KEY'] = 'bb79b746-999d-4bec-af22-04fddb05d087:49350e8b76fd8dda0fb9dd8442a9ccf5'
|
16 |
+
|
17 |
+
# Request body model
|
18 |
+
class StoryRequest(BaseModel):
|
19 |
+
mood: str
|
20 |
+
story_type: str
|
21 |
+
theme: str
|
22 |
+
num_scenes: int
|
23 |
+
txt: str
|
24 |
+
|
25 |
+
# Initialize the LLM
|
26 |
+
callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
|
27 |
+
|
28 |
+
tokenizer = AutoTokenizer.from_pretrained("tohur/natsumura-storytelling-rp-1.0-llama-3.1-8b")
|
29 |
+
model = AutoModelForCausalLM.from_pretrained("tohur/natsumura-storytelling-rp-1.0-llama-3.1-8b")
|
30 |
+
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=2000)
|
31 |
+
llm = HuggingFacePipeline(pipeline=pipe)
|
32 |
+
|
33 |
+
# Create a prompt template
|
34 |
+
# system = """You are a helpful and creative assistant that specializes in generating engaging and imaginative stories for kids.
|
35 |
+
# Based on the user's provided mood, preferred story type, theme, age, and desired story length of 500-600 words, create a unique and captivating story.
|
36 |
+
# Always start with Story Title then generate a single story and dont ask for any feedback at the end just sign off with a cute closing inviting the reader
|
37 |
+
# to create another adventure soon!
|
38 |
+
# """
|
39 |
+
|
40 |
+
system = """You are a helpful and creative assistant that specializes in generating engaging and imaginative short storie for kids.
|
41 |
+
Based on the user's provided mood, preferred story type, theme, age, and desired story length of 500-600 words, create a unique and captivating story.
|
42 |
+
Always start with Story Title then generate a single story.Storie begin on Page 1(also mention the all pages headings in bold) and end on Page 7.
|
43 |
+
Total pages in storie are seven each page have one short paragraph and dont ask for any feedback at the end just sign off with a cute closing inviting the reader
|
44 |
+
to create another adventure soon!
|
45 |
+
"""
|
46 |
+
|
47 |
+
prompt_template = ChatPromptTemplate.from_messages([("system", system), ("human", "{text}")])
|
48 |
+
|
49 |
+
# FastAPI endpoint to generate the story
|
50 |
+
@app.post("/generate_story/")
|
51 |
+
async def generate_story(story_request: StoryRequest):
|
52 |
+
story = f"""here are the inputs from user:
|
53 |
+
- **Mood:** {story_request.mood}
|
54 |
+
- **Story Type:** {story_request.story_type}
|
55 |
+
- **Theme:** {story_request.theme}
|
56 |
+
- **Details Provided:** {story_request.txt}
|
57 |
+
"""
|
58 |
+
|
59 |
+
final_prompt = prompt_template.format(text=story)
|
60 |
+
|
61 |
+
# Create the LLMChain
|
62 |
+
# chain = LLMChain(llm=llm, prompt=prompt_template)
|
63 |
+
chain = llm | prompt_template
|
64 |
+
|
65 |
+
# try:
|
66 |
+
# response = chain.invoke(final_prompt)
|
67 |
+
# return {"story": response}
|
68 |
+
# except Exception as e:
|
69 |
+
# raise HTTPException(status_code=500, detail=str(e))
|
70 |
+
response = chain.invoke(final_prompt)
|
71 |
+
|
72 |
+
if not response:
|
73 |
+
raise HTTPException(status_code=500, detail="Failed to generate the story")
|
74 |
+
|
75 |
+
images = []
|
76 |
+
for i in range(story_request.num_scenes):
|
77 |
+
# image_prompt = f"Generate an image for Scene {i+1} based on this story: Mood: {story_request.mood}, Story Type: {story_request.story_type}, Theme: {story_request.theme}. Story: {response}"
|
78 |
+
image_prompt = (
|
79 |
+
f"Generate an image for Scene {i+1}. "
|
80 |
+
f"This image should represent the details described in paragraph {i+1} of the story. "
|
81 |
+
f"Mood: {story_request.mood}, Story Type: {', '.join(story_request.story_type)}, Theme: {story_request.theme}. "
|
82 |
+
f"Story: {response} "
|
83 |
+
f"Focus on the key elements in paragraph {i+1}."
|
84 |
+
)
|
85 |
+
handler = fal_client.submit(
|
86 |
+
"fal-ai/flux/schnell",
|
87 |
+
arguments={
|
88 |
+
"prompt": image_prompt,
|
89 |
+
"num_images": 1,
|
90 |
+
"enable_safety_checker": True
|
91 |
+
},
|
92 |
+
)
|
93 |
+
result = handler.get()
|
94 |
+
image_url = result['images'][0]['url']
|
95 |
+
images.append(image_url)
|
96 |
+
|
97 |
+
return {
|
98 |
+
"story": response,
|
99 |
+
"images": images
|
100 |
+
}
|