Update app.py

app.py

@@ -1,16 +1,64 @@
 from flask import Flask, request, jsonify
-from flask_cors import CORS  # Add this import
+from flask_cors import CORS
 import logging
 import os
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+import torch
 
 app = Flask(__name__)
-CORS(app)  # Add this to handle CORS issues
+CORS(app)
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-@app.route('/generate_script', methods=['POST', 'GET'])  # Allow GET for testing
+# Initialize model variables
+model_name = "google/gemma-1b-it"
+tokenizer = None
+model = None
+generator = None
+
+def initialize_model():
+    global tokenizer, model, generator
+    
+    if tokenizer is None:
+        try:
+            # Get Hugging Face token from environment
+            hf_token = os.environ.get('HF_TOKEN')
+            
+            if not hf_token:
+                logger.error("HF_TOKEN environment variable not set!")
+                raise ValueError("Hugging Face token not configured")
+            
+            logger.info("Loading tokenizer...")
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_name,
+                token=hf_token
+            )
+            
+            logger.info("Loading model...")
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                token=hf_token
+            )
+            
+            logger.info("Creating pipeline...")
+            generator = pipeline(
+                "text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                device=0 if torch.cuda.is_available() else -1
+            )
+            
+            logger.info("Model loaded successfully!")
+            
+        except Exception as e:
+            logger.error(f"Error loading model: {str(e)}")
+            raise
+
+@app.route('/generate_script', methods=['POST', 'GET'])
 def generate_script():
     try:
         logger.info("Received request to /generate_script")
@@ -26,36 +74,58 @@ def generate_script():
         if not topics:
             return jsonify({"error": "No topics provided"}), 400
         
-        # For now, return a mock response to test the connection
-        mock_script = f"""VIRAL SCRIPT FOR: {topics[0]}
-
-Scene 1: Golden Hook
-- Visual: Eye-catching intro with text overlay "THIS WILL CHANGE EVERYTHING"
-- Voiceover: "What if I told you {topics[0].lower()} is about to revolutionize everything?"
+        # Initialize model (will load on first request)
+        initialize_model()
+        
+        # Create prompt for Gemma
+        prompt = f"""Create a viral 1-minute video script about {topics[0]}.
 
-Scene 2: Valuable Content
-- Visual: Demonstrating key points with graphics
-- Voiceover: "Here are 3 ways {topics[0].lower()} is changing the game..."
+REQUIREMENTS:
+- Duration: Exactly 1 minute
+- Include: Golden hook, valuable content, CTA
+- Virtual descriptions for each scene
+- Voiceover instructions for each scene
+- Make it engaging and viral-worthy
 
-Scene 3: CTA
-- Visual: Call-to-action screen with links
-- Voiceover: "Like and follow for more insights!"
+SCRIPT STRUCTURE:
 """
-
+        
+        # Generate script
+        logger.info("Generating script...")
+        result = generator(
+            prompt,
+            max_length=1024,
+            temperature=0.8,
+            do_sample=True,
+            top_p=0.9,
+            num_return_sequences=1
+        )
+        
+        generated_text = result[0]['generated_text']
+        
+        # Extract just the script part (remove the prompt)
+        script = generated_text.replace(prompt, "").strip()
+        
+        logger.info("Script generated successfully!")
+        
         return jsonify({
             "topic": topics[0],
-            "script": mock_script,
+            "script": script,
             "full_topics": topics,
             "status": "success"
         })
         
     except Exception as e:
-        logger.error(f"Error: {str(e)}")
+        logger.error(f"Error generating script: {str(e)}")
         return jsonify({"error": str(e)}), 500
 
 @app.route('/health', methods=['GET'])
 def health():
-    return jsonify({"status": "healthy", "message": "Server is running"})
+    return jsonify({
+        "status": "healthy", 
+        "message": "Server is running",
+        "model_loaded": tokenizer is not None
+    })
 
 @app.route('/test', methods=['GET'])
 def test():