Update main.py

main.py

@@ -7,16 +7,17 @@ from fastapi import FastAPI, File, UploadFile, Form, HTTPException
 
 # Removed 'requests' as we'll primarily use gradio_client for captioning
 # import requests
-import base64 # Still useful if you need base64 for anything else
+import base64 # Keep if you might need base64 for other purposes
 import os
 import random
+from typing import IO # Import IO for type hinting file-like objects
 
 # Import necessary classes from transformers
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Import the Gradio Client and handle_file
-from gradio_client import Client, handle_file # <--- Import handle_file
+from gradio_client import Client, handle_file
 
 from deep_translator import GoogleTranslator
 from deep_translator.exceptions import InvalidSourceOrTargetLanguage
@@ -24,150 +25,187 @@ from deep_translator.exceptions import InvalidSourceOrTargetLanguage
 
 app = FastAPI()
 
-# --- Hugging Face Model Setup (Local) ---
+# --- Hugging Face Language Model Setup (Local Inference) ---
 # Model name for TinyLlama 1.1B Chat (instruction-tuned version)
-# Or use "google/gemma-2b-it" if you got access and prefer its quality
+# Chosen for balance of quality and speed on CPU basic (faster than Gemma 2B, better than GPT-2 base)
+# If you get access to Gemma 2B-IT and prefer its quality (accepting slower speed),
+# change this to "google/gemma-2b-it" and use the generate_story_gemma function below.
 model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 tokenizer = None
 model = None
 
+# --- Hugging Face Gradio Space Client Setup (For Image Captioning) ---
 # Global Gradio Client for Captioning
 caption_client = None
-# The Space URL for the captioning API
+# The Space URL for the external captioning API
 CAPTION_SPACE_URL = "Makhinur/Image-to-Text-Salesforce-blip-image-captioning-base"
 
 
 # Function to load the language model and tokenizer
-def load_model():
+def load_language_model():
     global tokenizer, model
     print(f"Loading language model: {model_name}...")
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        # TinyLlama might not have an explicit pad token, setting it to EOS is common practice
+        if tokenizer.pad_token is None:
+             tokenizer.pad_token = tokenizer.eos_token
+
+        # Load model weights. Using float16 to reduce memory footprint on CPU.
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16, # Use float16 precision
+            # device_map="auto" # Not needed for single CPU
+        )
+        # model.to("cpu") # Explicitly ensure it's on CPU, although from_pretrained does this by default on CPU-only systems.
 
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    if tokenizer.pad_token is None:
-         tokenizer.pad_token = tokenizer.eos_token
-
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.float16, # Use float16 precision to save RAM
-        # device_map="auto" # Not needed for single CPU
-    )
-    # model.to("cpu") # Explicitly move if needed, though default is CPU
+        print(f"Language model {model_name} loaded successfully.")
+    except Exception as e:
+        print(f"Error loading language model {model_name}: {e}")
+        # Depending on requirements, you might want to exit or set model/tokenizer to None permanently
+        # For now, we let the app start, but subsequent generation calls will fail gracefully.
+        tokenizer = None
+        model = None
 
-    print(f"Language model {model_name} loaded successfully.")
 
-# Function to initialize the Gradio Client
+# Function to initialize the Gradio Client for the captioning Space
 def initialize_caption_client():
     global caption_client
     print(f"Initializing Gradio client for {CAPTION_SPACE_URL}...")
     try:
-        # Add your HF token here if the caption space becomes private or requires authentication
-        # HF_TOKEN = os.environ.get("HF_TOKEN") # Assume HF_TOKEN is set as a Space Secret
+        # If the target Gradio Space is private or requires authentication,
+        # uncomment the lines below and set your HF_TOKEN as a Space Secret.
+        # HF_TOKEN = os.environ.get("HF_TOKEN")
         # if HF_TOKEN:
         #     caption_client = Client(CAPTION_SPACE_URL, hf_token=HF_TOKEN)
         # else:
         #     caption_client = Client(CAPTION_SPACE_URL)
 
-        caption_client = Client(CAPTION_SPACE_URL) # Using public client for now
-
+        # Assuming the caption space is public and does not require a token
+        caption_client = Client(CAPTION_SPACE_URL)
         print("Gradio client initialized successfully.")
     except Exception as e:
-        print(f"Error initializing Gradio client: {e}")
+        print(f"Error initializing Gradio client for {CAPTION_SPACE_URL}: {e}")
+        # Set client to None so the endpoint can check and return an error
         caption_client = None
 
 
 # Load models and initialize clients when the app starts
 @app.on_event("startup")
 async def startup_event():
-    load_model()
+    # Load the language model (TinyLlama or Gemma)
+    load_language_model()
+    # Initialize the client for the captioning Space
     initialize_caption_client()
 
 
-# --- Image Captioning (Using gradio_client and handle_file) ---
-# Modify to accept UploadFile directly and use the gradio_client
+# --- Image Captioning Function (Using gradio_client) ---
 def generate_image_caption(image_file: UploadFile):
     """
     Generates a caption for the uploaded image using the external Gradio Space API.
-    Uses gradio_client and handle_file for correct file handling.
+    Reads the file content and uses handle_file for correct API input format.
     """
     if caption_client is None:
         error_msg = "Gradio caption client not initialized. Cannot generate caption."
-        print(error_msg)
-        return f"Error: {error_msg}"
+        print(error_msg) # Log the error server-side
+        return f"Error: {error_msg}" # Return an error string to the caller
 
     try:
         print(f"Calling caption API /predict for file {image_file.filename}...")
-        # Use handle_file() to prepare the file-like object for the API call
-        # Pass the file-like object (image_file.file) to handle_file
-        prepared_file = handle_file(image_file.file) # <--- Use handle_file here
 
-        # Call the predict method with the prepared file object
-        caption = caption_client.predict(img=prepared_file, api_name="/predict") # <--- Use prepared_file
+        # Read the content of the uploaded file into bytes
+        # It's important to seek(0) in case the file-like object has been read before
+        image_file.file.seek(0)
+        image_bytes = image_file.file.read()
+
+        # Use handle_file() with the byte content. This prepares the bytes
+        # into the format expected by the Gradio API (often base64).
+        prepared_input = handle_file(image_bytes)
+
+        # Call the predict method on the initialized client with the prepared input
+        caption = caption_client.predict(img=prepared_input, api_name="/predict")
 
         print(f"Caption generated: {caption}")
-        return caption
+        return caption # Return the successful caption string
     except Exception as e:
-        print(f"Error during caption generation API call: {e}")
-        # Provide more details in the error message for debugging
+        # Catch potential exceptions from gradio_client.predict (network, API error, etc.)
+        print(f"Error during caption generation API call: {e}") # Log the exception server-side
+        # Return an informative error string including the exception type and message
         return f"Error: Unable to generate caption from API. Details: {type(e).__name__}: {e}"
 
 
 # --- Language Model Story Generation Function ---
-# Use the appropriate function based on your chosen model (TinyLlama or Gemma)
-# This function name should match the model_name you've chosen.
-
+# This function uses the loaded TinyLlama model to generate the story.
 def generate_story_tinyllama(prompt_text: str, max_new_tokens: int = 300, temperature: float = 0.7, top_p: float = 0.9, top_k: int = 50) -> str:
     """
-    Generates text using the loaded TinyLlama model.
-    Applies the chat template.
+    Generates text using the loaded TinyLlama model based on the prompt.
+    Applies the TinyLlama Chat template.
     """
     if tokenizer is None or model is None:
-        raise RuntimeError("Language model and tokenizer not loaded. App startup failed?")
+        raise RuntimeError("Language model and tokenizer not loaded. Cannot generate story.")
 
+    # TinyLlama-Chat uses a chat template (similar to Llama/Gemma's instruction format)
     messages = [
         {"role": "user", "content": prompt_text}
+        # Add {"role": "system", "content": "Your system prompt here"} if needed
     ]
 
+    # Apply the chat template to format the prompt correctly for the model
     try:
          input_text = tokenizer.apply_chat_template(
              messages,
-             tokenize=False,
-             add_generation_prompt=True
+             tokenize=False, # Return as string before tokenizing
+             add_generation_prompt=True # Adds the assistant turn prompt token(s)
          )
     except AttributeError:
-         print("Warning: apply_chat_template not found. Using basic prompt formatting.")
-         input_text = f"<s>[INST] {prompt_text} [/INST]"
+         # Fallback for models that don't have a chat template defined
+         print("Warning: apply_chat_template not found for this tokenizer. Using basic prompt formatting.")
+         input_text = f"<s>[INST] {prompt_text} [/INST]" # Basic Llama/TinyLlama instruction format
+
+    # Encode the templated prompt into input IDs
+    # max_length should be within the model's context window (e.g., 4096 for TinyLlama)
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=1024) # Truncate if prompt is too long
 
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=1024)
+    # Ensure input tensors are on the same device as the model (CPU by default)
+    # inputs = {k: v.to(model.device) for k, v in inputs.items()} # Redundant if model is on CPU
 
+    # Generate new tokens based on the input prompt
     generate_ids = model.generate(
         inputs.input_ids,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k,
-        pad_token_id=tokenizer.pad_token_id,
+        max_new_tokens=max_new_tokens, # Maximum number of tokens to generate
+        do_sample=True, # Enable sampling for creative output
+        temperature=temperature, # Control randomness
+        top_p=top_p, # Control diversity (nucleus sampling)
+        top_k=top_k, # Control diversity (top-k sampling)
+        pad_token_id=tokenizer.pad_token_id, # Specify pad token for generation
+        # eos_token_id=tokenizer.eos_token_id # Optional: specify end-of-sequence token id to stop generation early
     )
 
+    # Decode the generated token IDs back into text
+    # Slice [0, inputs.input_ids.shape[-1]:] to get only the newly generated tokens
+    # skip_special_tokens=True removes tokens like <s>, </s>, <pad>
     generated_text = tokenizer.decode(generate_ids[0, inputs.input_ids.shape[-1]:], skip_special_tokens=True)
-    return generated_text.strip()
 
-# If using Gemma 2B instead of TinyLlama, use this function:
+    return generated_text.strip() # Return the generated text, stripped of leading/trailing whitespace
+
+# --- Optional: Gemma 2B Story Generation Function (if you prefer Gemma and get access) ---
 # def generate_story_gemma(prompt_text: str, max_new_tokens: int = 300, temperature: float = 0.7, top_p: float = 0.9, top_k: int = 50) -> str:
 #     """
-#     Generates text using the loaded Gemma model.
-#     Applies the Gemma-IT chat template.
+#     Generates text using the loaded Gemma model based on the prompt.
+#     Applies the Gemma-IT Chat template.
 #     """
 #     if tokenizer is None or model is None:
-#         raise RuntimeError("Language model and tokenizer not loaded. App startup failed?")
+#         raise RuntimeError("Language model and tokenizer not loaded. Cannot generate story.")
 
+#     # Gemma-IT uses a specific chat template
 #     messages = [
 #         {"role": "user", "content": prompt_text}
+#         # {"role": "system", "content": "Your system prompt here"} # Gemma also supports system prompts
 #     ]
 #     input_text = tokenizer.apply_chat_template(
 #         messages,
 #         tokenize=False,
-#         add_generation_prompt=True
+#         add_generation_prompt=True # Adds the assistant turn prompt token(s)
 #     )
 
 #     inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=1024)
@@ -186,9 +224,12 @@ def generate_story_tinyllama(prompt_text: str, max_new_tokens: int = 300, temper
 #     return generated_text.strip()
 
 
-# --- FastAPI Endpoint ---
+# --- FastAPI Endpoint for Story Generation ---
 @app.post("/generate-story/")
 async def generate_story_endpoint(image_file: UploadFile = File(...), language: str = Form(...)):
+    # No need to manually read image data here, generate_image_caption handles it
+
+    # Choose a random theme for the story prompt
     story_theme = random.choice([
         'an adventurous journey',
         'a mysterious encounter',
@@ -202,59 +243,91 @@ async def generate_story_endpoint(image_file: UploadFile = File(...), language:
         'a journey into the unknown'
     ])
 
-    # Get image caption using the gradio_client function
-    # Pass the UploadFile object directly
+    # Step 1: Get image caption using the external API via gradio_client
+    # Pass the UploadFile object directly to the captioning function
     caption = generate_image_caption(image_file)
-    if caption.startswith("Error"):
-        print(f"Caption generation failed: {caption}")
+
+    # Check if caption generation failed
+    if caption.startswith("Error:"): # Check if the returned string indicates an error
+        print(f"Caption generation failed: {caption}") # Log the error detail server-side
+        # Return a 500 Internal Server Error with the error message
         raise HTTPException(status_code=500, detail=caption)
 
-    # Construct the prompt for the language model
+    # Step 2: Construct the prompt for the language model
+    # We instruct the model to write a story based on the theme and incorporate the caption.
     prompt_text = f"Write an attractive story of around 300 words about {story_theme}. Incorporate the following details from an image description into the story: {caption}\n\nStory:"
 
-    # Generate the story using the appropriate function (adjust if using Gemma)
+    # Step 3: Generate the story using the local language model
     try:
-        story = generate_story_tinyllama( # <--- Make sure this matches your chosen model function
+        # Call the appropriate story generation function (TinyLlama in this case)
+        story = generate_story_tinyllama(
             prompt_text,
-            max_new_tokens=300,
-            temperature=0.7,
+            max_new_tokens=300, # Aim for ~300 new tokens
+            temperature=0.7,    # Standard creative sampling parameters
             top_p=0.9,
             top_k=50
         )
-        story = story.strip()
+        story = story.strip() # Basic cleanup of potential extra whitespace
 
+    except RuntimeError as e:
+        # Catch errors specifically from model not being loaded
+        print(f"Language model not loaded error: {e}")
+        raise HTTPException(status_code=503, detail=f"Language model not available: {e}")
     except Exception as e:
-        print(f"Story generation failed: {e}")
-        raise HTTPException(status_code=500, detail=f"Story generation failed: {e}. Please check Space logs for details.")
+        # Catch other potential errors during generation
+        print(f"Story generation failed: {e}") # Log the exception server-side
+        raise HTTPException(status_code=500, detail=f"Story generation failed: {type(e).__name__}: {e}. Please check Space logs for details.")
 
 
-    # Translate the story
-    if language.lower() != "english":
+    # Step 4: Translate the story if the target language is not English
+    if language and language.lower() != "english":
         try:
+            # Use GoogleTranslator with specified source and target languages
             translator = GoogleTranslator(source='english', target=language.lower())
             translated_story = translator.translate(story)
 
-            if translated_story is None:
-                 print(f"Translation returned None for language: {language}")
+            # Check if translation was successful or returned None
+            if translated_story is None or translated_story == "":
+                 print(f"Translation returned None or empty string for language: {language}")
+                 # If translation fails, return the English story with a warning message
                  return {"story": story + "\n\n(Note: Automatic translation to your requested language failed.)"}
 
-            story = translated_story
+            story = translated_story # Use the translated story
 
         except InvalidSourceOrTargetLanguage:
              print(f"Invalid target language requested: {language}")
              raise HTTPException(status_code=400, detail=f"Invalid target language: {language}")
         except Exception as e:
-             print(f"Translation failed for language {language}: {e}")
-             raise HTTPException(status_code=500, detail=f"Translation failed: {e}")
+             # Catch other potential translation errors (network, API issues, etc.)
+             print(f"Translation failed for language {language}: {e}") # Log server-side
+             raise HTTPException(status_code=500, detail=f"Translation failed: {type(e).__name__}: {e}")
+
 
+    # Step 5: Return the final generated (and potentially translated) story
     return {"story": story}
 
-# --- Optional: HTML form for testing (Needs templates dir and index.html) ---
+# --- Optional: Serve a simple HTML form for testing ---
+# To enable this, create a 'templates' directory and an 'index.html' file inside it.
+# Also uncomment the imports at the top related to HTMLResponse, StaticFiles, Jinja2Templates, Request.
 # from fastapi import Request
 # from fastapi.templating import Jinja2Templates
 # from fastapi.staticfiles import StaticFiles
 # templates = Jinja2Templates(directory="templates")
 # app.mount("/static", StaticFiles(directory="static"), name="static")
+
 # @app.get("/", response_class=HTMLResponse)
 # async def read_root(request: Request):
+#     # Example index.html structure for a simple form:
+#     # <!DOCTYPE html>
+#     # <html>
+#     # <head><title>Story Generator</title></head>
+#     # <body>
+#     #   <h1>Generate a Story from an Image</h1>
+#     #   <form action="/generate-story/" method="post" enctype="multipart/form-data">
+#     #     <input type="file" name="image_file" accept="image/*" required><br><br>
+#     #     Target Language (e.g., english, french, spanish): <input type="text" name="language" value="english"><br><br>
+#     #     <button type="submit">Generate Story</button>
+#     #   </form>
+#     # </body>
+#     # </html>
 #     return templates.TemplateResponse("index.html", {"request": request})