Update main.py

main.py

@@ -5,12 +5,12 @@ from fastapi import FastAPI, File, UploadFile, Form, HTTPException
 # from fastapi.templating import Jinja2Templates
 # from fastapi.responses import FileResponse
 
-# Removed 'requests' as we'll primarily use gradio_client for captioning
+# Removed 'requests'
 # import requests
-import base64 # Keep if you might need base64 for other purposes
+import base64 # Keep if needed elsewhere
 import os
 import random
-from typing import IO # Import IO for type hinting file-like objects
+from typing import IO
 
 # Import necessary classes from transformers
 import torch
@@ -19,6 +19,11 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 # Import the Gradio Client and handle_file
 from gradio_client import Client, handle_file
 
+# Import necessary modules for temporary file handling
+import tempfile
+import shutil # Useful for potential cleanup, although os.remove is sufficient here
+
+
 from deep_translator import GoogleTranslator
 from deep_translator.exceptions import InvalidSourceOrTargetLanguage
 
@@ -26,10 +31,7 @@ from deep_translator.exceptions import InvalidSourceOrTargetLanguage
 app = FastAPI()
 
 # --- Hugging Face Language Model Setup (Local Inference) ---
-# Model name for TinyLlama 1.1B Chat (instruction-tuned version)
-# 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.
+# Using TinyLlama 1.1B Chat as the example
 model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 tokenizer = None
 model = None
@@ -47,23 +49,17 @@ def load_language_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.
 
         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
 
@@ -73,69 +69,85 @@ def initialize_caption_client():
     global caption_client
     print(f"Initializing Gradio client for {CAPTION_SPACE_URL}...")
     try:
-        # If the target Gradio Space is private or requires authentication,
-        # uncomment the lines below and set your HF_TOKEN as a Space Secret.
+        # If the target Space is private or requires authentication,
+        # uncomment and set 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)
 
-        # Assuming the caption space is public and does not require a token
+        # Assuming the caption space is public
         caption_client = Client(CAPTION_SPACE_URL)
         print("Gradio client initialized successfully.")
     except Exception as 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 the language model (TinyLlama or Gemma)
     load_language_model()
-    # Initialize the client for the captioning Space
     initialize_caption_client()
 
 
-# --- Image Captioning Function (Using gradio_client) ---
+# --- Image Captioning Function (Using gradio_client and temporary file) ---
 def generate_image_caption(image_file: UploadFile):
     """
     Generates a caption for the uploaded image using the external Gradio Space API.
-    Reads the file content and uses handle_file for correct API input format.
+    Saves the uploaded file to a temporary file and uses its path with handle_file.
     """
     if caption_client is None:
         error_msg = "Gradio caption client not initialized. Cannot generate caption."
-        print(error_msg) # Log the error server-side
-        return f"Error: {error_msg}" # Return an error string to the caller
+        print(error_msg)
+        return f"Error: {error_msg}"
+
+    temp_file_path = None # Initialize temporary file path variable
 
     try:
         print(f"Calling caption API /predict for file {image_file.filename}...")
 
-        # 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)
+        # Read the content of the uploaded file
+        image_file.file.seek(0) # Ensure pointer is at the beginning
         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)
+        # Create a temporary file and write the bytes to it
+        # use delete=False so the file isn't deleted automatically when closed
+        # Use suffix to hint at the file type
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(image_file.filename)[1] or '.jpg')
+        temp_file.write(image_bytes)
+        temp_file.close() # Close the handle so gradio_client can access it
+        temp_file_path = temp_file.name # Get the actual path string
+
+        print(f"Saved uploaded file temporarily to: {temp_file_path}")
+
+        # Use handle_file() with the path to the temporary file (a string)
+        # This aligns with the documentation "img filepath Required" and example handle_file('...')
+        prepared_input = handle_file(temp_file_path)
 
         # 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 the successful caption string
+
     except Exception as e:
-        # 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
+        # Catch potential exceptions during the process
+        print(f"Error during caption generation API call: {e}")
         return f"Error: Unable to generate caption from API. Details: {type(e).__name__}: {e}"
 
+    finally:
+        # Clean up the temporary file
+        if temp_file_path and os.path.exists(temp_file_path):
+            print(f"Cleaning up temporary file: {temp_file_path}")
+            try:
+                os.remove(temp_file_path)
+            except OSError as e:
+                 print(f"Error removing temporary file {temp_file_path}: {e}") # Log cleanup errors
+
 
-# --- Language Model Story Generation Function ---
-# This function uses the loaded TinyLlama model to generate the story.
+# --- Language Model Story Generation Function (TinyLlama) ---
 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 based on the prompt.
@@ -144,117 +156,56 @@ def generate_story_tinyllama(prompt_text: str, max_new_tokens: int = 300, temper
     if tokenizer is None or model is None:
         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, # Return as string before tokenizing
-             add_generation_prompt=True # Adds the assistant turn prompt token(s)
+             tokenize=False,
+             add_generation_prompt=True
          )
     except AttributeError:
-         # 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
+         print("Warning: apply_chat_template not found. Using basic prompt formatting.")
+         input_text = f"<s>[INST] {prompt_text} [/INST]"
 
-    # 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, # 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
+        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,
     )
 
-    # 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() # 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 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. 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 # Adds the assistant turn prompt token(s)
-#     )
-
-#     inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=1024)
-
-#     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,
-#     )
-
-#     generated_text = tokenizer.decode(generate_ids[0, inputs.input_ids.shape[-1]:], skip_special_tokens=True)
-#     return generated_text.strip()
-
+    return generated_text.strip()
 
 # --- 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
+    # Choose a random theme
     story_theme = random.choice([
-        'an adventurous journey',
-        'a mysterious encounter',
-        'a heroic quest',
-        'a magical adventure',
-        'a thrilling escape',
-        'an unexpected discovery',
-        'a dangerous mission',
-        'a romantic escapade',
-        'an epic battle',
+        'an adventurous journey', 'a mysterious encounter', 'a heroic quest',
+        'a magical adventure', 'a thrilling escape', 'an unexpected discovery',
+        'a dangerous mission', 'a romantic escapade', 'an epic battle',
         'a journey into the unknown'
     ])
 
     # Step 1: Get image caption using the external API via gradio_client
-    # Pass the UploadFile object directly to the captioning function
+    # Pass the UploadFile object to the captioning function
     caption = generate_image_caption(image_file)
 
     # 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
+    if caption.startswith("Error:"):
+        print(f"Caption generation failed: {caption}")
         raise HTTPException(status_code=500, detail=caption)
 
     # 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:"
 
     # Step 3: Generate the story using the local language model
@@ -262,72 +213,50 @@ async def generate_story_endpoint(image_file: UploadFile = File(...), language:
         # Call the appropriate story generation function (TinyLlama in this case)
         story = generate_story_tinyllama(
             prompt_text,
-            max_new_tokens=300, # Aim for ~300 new tokens
-            temperature=0.7,    # Standard creative sampling parameters
+            max_new_tokens=300,
+            temperature=0.7,
             top_p=0.9,
             top_k=50
         )
-        story = story.strip() # Basic cleanup of potential extra whitespace
+        story = story.strip()
 
     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:
-        # Catch other potential errors during generation
-        print(f"Story generation failed: {e}") # Log the exception server-side
+        print(f"Story generation failed: {e}")
         raise HTTPException(status_code=500, detail=f"Story generation failed: {type(e).__name__}: {e}. Please check Space logs for details.")
 
 
     # 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)
 
-            # 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 # Use the translated story
+            story = 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:
-             # Catch other potential translation errors (network, API issues, etc.)
-             print(f"Translation failed for language {language}: {e}") # Log server-side
+             print(f"Translation failed for language {language}: {e}")
              raise HTTPException(status_code=500, detail=f"Translation failed: {type(e).__name__}: {e}")
 
-
-    # Step 5: Return the final generated (and potentially translated) story
+    # Step 5: Return the final story
     return {"story": story}
 
 # --- 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.
+# Requires uncommenting imports above and creating 'templates/index.html'
 # 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})