Update main.py

main.py

@@ -44,6 +44,13 @@ except Exception as e:
 
 bucket = storage.bucket()
 
+
+# Helper: Upload a local file to Firebase Storage and return its public URL
+def upload_to_storage(local_path, destination_blob_name):
+    blob = bucket.blob(destination_blob_name)
+    blob.upload_from_filename(local_path)
+    return blob.public_url
+    
 # Gemini API initialization
 api_key = os.environ['Gemini']
 def configure_gemini():
@@ -161,89 +168,240 @@ def get_user_profile():
 # -----------------------
 # Content
 # -----------------------
+# ---------- Story Generation Endpoint ----------
 
-# ---------- Video Generation Endpoint ----------
-
-@app.route('/api/video/generate', methods=['POST'])
-def generate_video():
+@app.route('/api/story/generate', methods=['POST'])
+def generate_story_endpoint():
     try:
+        # --- Authentication ---
         auth_header = request.headers.get('Authorization', '')
         if not auth_header.startswith('Bearer '):
-            return jsonify({'error': 'Authorization header missing or malformed'}), 401
+            return jsonify({'error': 'Missing or invalid token'}), 401
         token = auth_header.split(' ')[1]
         uid = verify_token(token)
         if not uid:
-            return jsonify({'error': 'Invalid token'}), 401
-
-        # Get user data and check credits
-        user_ref = db.reference(f'users/{uid}')
-        user_data = user_ref.get()
-        if not user_data or user_data.get('credits', 0) < 1:
-            return jsonify({'error': 'Insufficient credits'}), 400
+            return jsonify({'error': 'Invalid or expired token'}), 401
 
-        # Deduct one credit
-        new_credits = user_data.get('credits', 0) - 5
-        user_ref.update({'credits': new_credits})
+        # --- Read Request Data ---
+        data = request.get_json()
+        input_type = data.get('input_type', 'text')  # "text", "wiki", "bible", "youtube", or "dataframe"
+        prompt = data.get('prompt')
+        custom_prompt = data.get('custom_prompt', "")  # Optional extra prompt
+        story_type = data.get('story_type', 'free_form')
+        style = data.get('style', 'whimsical')
+        voice_model = data.get('voice_model', 'aura-asteria-en')
+        image_model = data.get('image_model', 'hf')
+        audio_model = data.get('audio_model', 'deepgram')
 
-        # Get prompt and generate video (simulation using Gemini)
-        req_data = request.get_json()
-        prompt = req_data.get('prompt', '')
         if not prompt:
             return jsonify({'error': 'Prompt is required'}), 400
 
-        generated_video = ""
+        # Combine prompt and custom_prompt (if provided)
+        combined_prompt = prompt + (" " + custom_prompt if custom_prompt else "")
+
+        # --- Select the Appropriate Story Generation Function ---
+        story_gen_start = time.time()
+        full_story = None
+
+        if input_type == "text":
+            from stories import generate_story_from_text
+            full_story = generate_story_from_text(combined_prompt, story_type)
+        elif input_type == "wiki":
+            wiki_url = data.get("wiki_url")
+            if not wiki_url:
+                return jsonify({'error': 'wiki_url is required for input_type "wiki"'}), 400
+            from stories import generate_story_from_wiki
+            full_story = generate_story_from_wiki(wiki_url, story_type)
+        elif input_type == "bible":
+            bible_reference = data.get("bible_reference")
+            if not bible_reference:
+                return jsonify({'error': 'bible_reference is required for input_type "bible"'}), 400
+            from stories import generate_story_from_bible
+            full_story = generate_story_from_bible(bible_reference, story_type)
+        elif input_type == "youtube":
+            youtube_url = data.get("youtube_url")
+            if not youtube_url:
+                return jsonify({'error': 'youtube_url is required for input_type "youtube"'}), 400
+            from stories import generate_story_from_youtube
+            full_story = generate_story_from_youtube(youtube_url, story_type)
+        elif input_type == "dataframe":
+            # Expecting dataframe data as JSON (list of dicts)
+            df_data = data.get("data")
+            if not df_data:
+                return jsonify({'error': 'Data for dataframe input_type is required'}), 400
+            df = pd.DataFrame(df_data)
+            from stories import generate_story_from_dataframe
+            full_story = generate_story_from_dataframe(df, story_type)
+        else:
+            return jsonify({'error': 'Unsupported input_type'}), 400
+
+        story_gen_end = time.time()
+        story_generation_time = story_gen_end - story_gen_start
+
+        if not full_story:
+            return jsonify({'error': 'Story generation failed'}), 500
+
+        # --- Split the Story into 5 Sections ---
+        sections_raw = [s.strip() for s in full_story.split("[break]") if s.strip()]
+        if len(sections_raw) < 5:
+            sections_raw += ["(Placeholder section)"] * (5 - len(sections_raw))
+        elif len(sections_raw) > 5:
+            sections_raw = sections_raw[:5]
+
+        sections = []
+        image_generation_times = []
+        audio_generation_times = []
+
+        # Import generation functions from your modules
+        from image_gen import generate_image_with_retry  # image generation function
+        from audio_gen import generate_audio              # audio generation function
+
+        # Process each section
+        for section in sections_raw:
+            # --- Image Generation ---
+            # Extract an image prompt between angle brackets; otherwise, fallback to the first 100 characters.
+            img_prompt_match = re.search(r"<(.*?)>", section)
+            img_prompt = img_prompt_match.group(1).strip() if img_prompt_match else section[:100]
+
+            image_start = time.time()
+            image_obj, _ = generate_image_with_retry(img_prompt, style, model=image_model)
+            image_end = time.time()
+            image_generation_times.append(image_end - image_start)
+
+            # Save image locally and upload it.
+            image_filename = f"/tmp/{uuid.uuid4().hex}.jpg"
+            image_obj.save(image_filename, format="JPEG")
+            image_blob_name = f"stories/{uid}/{uuid.uuid4().hex}.jpg"
+            image_url = upload_to_storage(image_filename, image_blob_name)
+
+            # --- Audio Generation ---
+            audio_start = time.time()
+            audio_file_path = generate_audio(section, voice_model, audio_model=audio_model)
+            audio_end = time.time()
+            audio_generation_times.append(audio_end - audio_start)
+
+            audio_blob_name = f"stories/{uid}/{uuid.uuid4().hex}.mp3"
+            audio_url = upload_to_storage(audio_file_path, audio_blob_name)
+
+            sections.append({
+                "section_text": section,
+                "image_url": image_url,
+                "audio_url": audio_url
+            })
+
+            # Clean up temporary files
+            os.remove(image_filename)
+            os.remove(audio_file_path)
+
+        # --- Store the Story Record in Firebase Realtime Database ---
+        story_id = str(uuid.uuid4())
+        story_ref = db.reference(f"stories/{story_id}")
+        story_record = {
+            "uid": uid,
+            "full_story": full_story,
+            "sections": sections,
+            "generation_times": {
+                "story_generation_time": story_generation_time,
+                "image_generation_times": image_generation_times,
+                "audio_generation_times": audio_generation_times
+            },
+            "created_at": datetime.utcnow().isoformat(),
+            "input_type": input_type,
+            "story_type": story_type
+        }
+        story_ref.set(story_record)
+
+        preview = sections[0] if sections else {}
 
         return jsonify({
-            'success': True,
-            'video': generated_video,
-            'remaining_credits': new_credits
+            "story_id": story_id,
+            "full_story": full_story,
+            "preview": preview,
+            "sections": sections,
+            "generation_times": story_record["generation_times"]
         })
+
     except Exception as e:
         return jsonify({'error': str(e)}), 500
 
-# ---------- Story Generation Endpoint ----------
-
-@app.route('/api/story/generate', methods=['POST'])
-def generate_story():
+# ---------- Video Generation Endpoint ----------
+@app.route('/api/video/generate', methods=['POST'])
+def generate_video_endpoint():
     try:
+        # Verify user token
         auth_header = request.headers.get('Authorization', '')
         if not auth_header.startswith('Bearer '):
-            return jsonify({'error': 'Authorization header missing or malformed'}), 401
+            return jsonify({'error': 'Missing or invalid token'}), 401
         token = auth_header.split(' ')[1]
         uid = verify_token(token)
         if not uid:
-            return jsonify({'error': 'Invalid token'}), 401
-
-        # Get user data and check credits
-        user_ref = db.reference(f'users/{uid}')
-        user_data = user_ref.get()
-        if not user_data or user_data.get('credits', 0) < 1:
-            return jsonify({'error': 'Insufficient credits'}), 400
-
-        # Deduct one credit
-        new_credits = user_data.get('credits', 0) - 5
-        user_ref.update({'credits': new_credits})
-
-        # Get prompt and generate video (simulation using Gemini)
-        req_data = request.get_json()
-        prompt = req_data.get('prompt', '')
-        if not prompt:
-            return jsonify({'error': 'Prompt is required'}), 400
-
-        model = configure_gemini()
-        # Call to Gemini to generate video content (this is a simulation)
-        response = "function(s) that does something"
-        # For demonstration, we assume the response contains a 'result' field
-        generated_story = response.result if hasattr(response, 'result') else "story generated based on parameters."
+            return jsonify({'error': 'Invalid or expired token'}), 401
 
-        return jsonify({
-            'success': True,
-            'story': generated_story,
-            'remaining_credits': new_credits
-        })
+        data = request.get_json()
+        story_id = data.get('story_id')
+        if not story_id:
+            return jsonify({'error': 'story_id is required'}), 400
+
+        # Retrieve the story record from the database.
+        story_ref = db.reference(f"stories/{story_id}")
+        story_data = story_ref.get()
+        if not story_data:
+            return jsonify({'error': 'Story not found'}), 404
+
+        sections = story_data.get("sections", [])
+        if not sections:
+            return jsonify({'error': 'No sections found in the story'}), 404
+
+        # Download the image and audio files from their URLs.
+        image_files = []
+        audio_files = []
+        for section in sections:
+            image_url = section.get("image_url")
+            audio_url = section.get("audio_url")
+            # Download image
+            img_resp = requests.get(image_url)
+            if img_resp.status_code == 200:
+                img_path = f"/tmp/{uuid.uuid4().hex}.jpg"
+                with open(img_path, "wb") as f:
+                    f.write(img_resp.content)
+                image_files.append(img_path)
+            # Download audio
+            aud_resp = requests.get(audio_url)
+            if aud_resp.status_code == 200:
+                aud_path = f"/tmp/{uuid.uuid4().hex}.mp3"
+                with open(aud_path, "wb") as f:
+                    f.write(aud_resp.content)
+                audio_files.append(aud_path)
+
+        if not image_files:
+            return jsonify({'error': 'No images available for video generation'}), 500
+
+        # Generate video using the video generation function (from video_gen.py)
+        from video_gen import create_video  # from your video_gen.py
+        video_output_path = f"/tmp/{uuid.uuid4().hex}.mp4"
+        video_file = create_video(image_files, audio_files, output_path=video_output_path)
+        if not video_file:
+            return jsonify({'error': 'Video generation failed'}), 500
+
+        # Upload the video to Firebase Storage.
+        video_blob_name = f"stories/{uid}/{uuid.uuid4().hex}.mp4"
+        video_url = upload_to_storage(video_file, video_blob_name)
+
+        # Optionally, update the story record with the video URL.
+        story_ref.update({"video_url": video_url})
+
+        # Clean up temporary local files.
+        for f in image_files:
+            os.remove(f)
+        for f in audio_files:
+            os.remove(f)
+        os.remove(video_file)
+
+        return jsonify({"video_url": video_url})
     except Exception as e:
         return jsonify({'error': str(e)}), 500
 
+
 # ---------- Audio Generation Endpoint ----------
 @app.route('/api/audio/generate', methods=['POST'])
 def generate_audio_edit():
@@ -272,15 +430,11 @@ def generate_audio_edit():
         if not prompt:
             return jsonify({'error': 'Prompt is required'}), 400
 
-        model = configure_gemini()
-        # Call to Gemini to generate video content (this is a simulation)
-        response = model.generate(prompt=prompt)
-        # For demonstration, we assume the response contains a 'result' field
-        generated_video = response.result if hasattr(response, 'result') else "Video content generated based on prompt."
+        
 
         return jsonify({
             'success': True,
-            'video': generated_video,
+            'audio': generated_audio,
             'remaining_credits': new_credits
         })
     except Exception as e:
@@ -317,59 +471,13 @@ def generate_image_edit():
 
         model = configure_gemini()
         # Call to Gemini to generate video content (this is a simulation)
-        response = model.generate(prompt=prompt)
-        # For demonstration, we assume the response contains a 'result' field
-        generated_video = response.result if hasattr(response, 'result') else "Video content generated based on prompt."
-
+        generated_image = ""
         return jsonify({
-            'success': True,
-            'video': generated_video,
-            'remaining_credits': new_credits
+            'audio': '.mp3'
         })
     except Exception as e:
         return jsonify({'error': str(e)}), 500
 
-# ---------- Preview Generation Endpoint ----------
-@app.route('/api/preview/generate', methods=['POST'])
-def generate_preview():
-    try:
-        auth_header = request.headers.get('Authorization', '')
-        if not auth_header.startswith('Bearer '):
-            return jsonify({'error': 'Authorization header missing or malformed'}), 401
-        token = auth_header.split(' ')[1]
-        uid = verify_token(token)
-        if not uid:
-            return jsonify({'error': 'Invalid token'}), 401
-
-        # Get user data and check credits
-        user_ref = db.reference(f'users/{uid}')
-        user_data = user_ref.get()
-        if not user_data or user_data.get('credits', 0) < 1:
-            return jsonify({'error': 'Insufficient credits'}), 400
-
-        # Deduct one credit
-        new_credits = user_data.get('credits', 0) - 1
-        user_ref.update({'credits': new_credits})
-
-        # Get prompt and generate video (simulation using Gemini)
-        req_data = request.get_json()
-        prompt = req_data.get('prompt', '')
-        if not prompt:
-            return jsonify({'error': 'Prompt is required'}), 400
-
-        model = configure_gemini()
-        # Call to Gemini to generate video content (this is a simulation)
-        response = model.generate(prompt=prompt)
-        # For demonstration, we assume the response contains a 'result' field
-        generated_video = response.result if hasattr(response, 'result') else "Video content generated based on prompt."
-
-        return jsonify({
-            'success': True,
-            'video': generated_video,
-            'remaining_credits': new_credits
-        })
-    except Exception as e:
-        return jsonify({'error': str(e)}), 500
 
 # ---------- Credit Request Endpoints ----------