Successfully creating the application that allows for story generation based on the given parameters and image generation (lagging) for now due to overloaded image size, will further optimize

app.py

@@ -84,13 +84,12 @@ class VisualNovelGenerator:
             layout="wide",
             initial_sidebar_state="expanded"
         )
-
-    def image_api_request(prompt: str) -> str:
+    def image_api_request(self, prompt: str) -> str:
         """Send a request to generate an image based on the prompt"""
-        image_api_key = st.secrets.get("IMG_API_KEY") or st.session_state.get("img_api_key")
+        image_api_key = st.secrets.get("IMG_API_KEY") or st.session_state.get("IMG_API_KEY")
         client = InferenceClient("black-forest-labs/FLUX.1-dev", token=image_api_key)
         try:
-            image = client.text_to_image(prompt)
+            image = client.text_to_image(prompt, target_size=(512, 512))    
             return image
         except Exception as e:
             st.error(f"Image Generation Error: {str(e)}")
@@ -180,7 +179,6 @@ class VisualNovelGenerator:
             for i in range(num_characters):
                 with st.expander(f"Character {i+1}"):
                     char_name = st.text_input("Name", key=f"char_name_{i}")
-
                     # Let's generate character details using AI for each character
                     if char_name:
                         prompt = f"""Return a JSON object for a character named {char_name} in this format exactly:
@@ -206,6 +204,28 @@ class VisualNovelGenerator:
                                     char_info = json.loads(json_str)
                                     characters[char_name] = char_info
                                     st.success(f"Character {char_name} details generated!")
+
+                                # Remove any non-JSON text before and after the JSON object
+                                json_start = response.find("{")
+                                json_end = response.rfind("}") + 1
+                                if json_start >= 0 and json_end > json_start:
+                                    json_str = response[json_start:json_end]
+                                    char_info = json.loads(json_str)
+                                    characters[char_name] = char_info
+                                    st.success(f"Character {char_name} details generated!")
+
+                                # Generate character portrait
+                                portrait_prompt = f"Visual novel style portrait of {char_name}: full face, shoulders up, anime style, high quality render"
+                                # Generate and save character portrait
+                                portrait_filename = f"portrait_{char_name.lower().replace(' ', '_')}.png"
+                                if not Path(portrait_filename).exists():
+                                    portrait = self.image_api_request(portrait_prompt)
+                                    if portrait:
+                                        portrait.save(portrait_filename)
+                                        st.image(portrait_filename, caption=f"Portrait of {char_name}")
+                                else:
+                                    st.image(portrait_filename, caption=f"Portrait of {char_name}")
+
                             except:
                                 st.error(f"Error generating details for {char_name}")
 
@@ -324,6 +344,18 @@ class VisualNovelGenerator:
                 "choices": story_data["dialogue_options"],
                 "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
             })
+
+            # Generate images for the opening scene and characters
+            try:
+                # Generate scene image
+                scene_prompt = f"Visual novel style, {setting} setting: {story_data['opening_scene']}"
+                scene_image = self.image_api_request(scene_prompt)
+                if scene_image:
+                    scene_image.save("temp_scene.png")
+                    st.image("temp_scene.png", caption="Opening Scene")
+
+            except Exception as e:
+                st.warning(f"Image generation failed: {str(e)}")
             
         except Exception as e:
             st.error(f"Error generating story: {str(e)}")

requirements.txt

@@ -1 +1,6 @@
 streamlit
+huggingface_hub
+re
+requests
+json
+openai