Update main.py

main.py

@@ -1,101 +1,219 @@
 from fastapi import FastAPI, File, UploadFile, Form, HTTPException
-from fastapi.responses import HTMLResponse
-from fastapi.staticfiles import StaticFiles
-from fastapi.templating import Jinja2Templates
-from fastapi.responses import FileResponse
+# Keep these if you use them elsewhere in your app (HTML, static files)
+# from fastapi.responses import HTMLResponse
+# from fastapi.staticfiles import StaticFiles
+# from fastapi.templating import Jinja2Templates
+# from fastapi.responses import FileResponse
+
 import requests
 import base64
-from typing import Iterator
 import os
-from text_generation import Client
+import random
+
+# Import necessary classes from transformers
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig # Added BitsAndBytesConfig in case you ever need quantization
+
+
 from deep_translator import GoogleTranslator
+from deep_translator.exceptions import InvalidSourceOrTargetLanguage
+
 
 app = FastAPI()
 
-model_id = 'codellama/CodeLlama-34b-Instruct-hf'
-
-API_URL = "https://api-inference.huggingface.co/models/" + model_id
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
-
-client = Client(
-    API_URL,
-    headers={"Authorization": f"Bearer {HF_TOKEN}"},
-)
-EOS_STRING = "</s>"
-EOT_STRING = "<EOT>"
-
-
-def get_prompt(message: str, chat_history: list[tuple[str, str]],
-               system_prompt: str) -> str:
-    texts = [f'<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n']
-    do_strip = False
-    for user_input, response in chat_history:
-        user_input = user_input.strip() if do_strip else user_input
-        do_strip = True
-        texts.append(f'{user_input} [/INST] {response.strip()} </s><s>[INST] ')
-    message = message.strip() if do_strip else message
-    texts.append(f'{message} [/INST]')
-    return ''.join(texts)
-
-
-def run(message: str,
-        chat_history: list[tuple[str, str]],
-        system_prompt: str,
-        max_new_tokens: int = 1024,
-        temperature: float = 0.1,
-        top_p: float = 0.9,
-        top_k: int = 50) -> Iterator[str]:
-    prompt = get_prompt(message, chat_history, system_prompt)
-
-    generate_kwargs = dict(
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
+# --- Hugging Face Model Setup (Local) ---
+# Model name for Gemma 2B Instruction-Tuned
+# This version is trained to follow instructions, ideal for your task.
+model_name = "google/gemma-2b-it"
+tokenizer = None
+model = None
+
+# Function to load the model and tokenizer
+def load_model():
+    global tokenizer, model
+    print(f"Loading model: {model_name}...")
+
+    # Load tokenizer
+    # trust_remote_code=True might be needed for some newer models/features,
+    # but standard Gemma usually works without it. Let's omit it for security unless necessary.
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    # Load model - Gemma can be loaded in float16 to save RAM
+    # On CPU, float16 performance can vary, but it reduces memory bandwidth
+    # which can sometimes help. 16GB RAM is plenty for Gemma 2B float16 (~2GB).
+    # We don't need quantization (load_in_8bit/4bit) for Gemma 2B with 16GB RAM,
+    # but it's an option for larger models or less RAM.
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float16, # Use float16 precision
+        # device_map="auto" # Not strictly needed for single CPU inference
     )
-    stream = client.generate_stream(prompt, **generate_kwargs)
-    output = ""
-    for response in stream:
-        if any([end_token in response.token.text for end_token in [EOS_STRING, EOT_STRING]]):
-            yield output
-            output = ""
-        else:
-            output += response.token.text
+    # model.to("cpu") # Explicitly ensure it's on CPU, although from_pretrained on CPU does this.
 
+    print(f"Model {model_name} loaded successfully.")
 
+# Load the model when the app starts
+@app.on_event("startup")
+async def startup_event():
+    load_model()
+
+# --- Image Captioning (External API - Keep) ---
+# Keep this as it is, it uses an external service
 def generate_image_caption(image_data):
-    image_base64 = base64.b64encode(image_data).decode('utf-8')
-    payload = {"data": ["data:image/jpeg;base64," + image_base64]}
+    payload = {"data": ["data:image/jpeg;base64," + base64.b64encode(image_data).decode('utf-8')]}
+    # Use the correct URL for the captioning API. This is the one from your original code.
+    # Ensure it's stable or replace if needed.
     response = requests.post("https://makhinur-image-to-text-salesforce-blip-image-cap-c0a9076.hf.space/run/predict", json=payload)
     if response.status_code == 200:
-        caption = response.json()["data"][0]
-        return caption
+        try:
+            result = response.json()
+            caption = result.get("data", ["Error: Unexpected API response format"])[0]
+            return caption
+        except Exception as e:
+            return f"Error: Failed to parse caption API response: {e}"
     else:
-        return "Error: Unable to generate caption"
+        return f"Error: Caption API returned status code {response.status_code}: {response.text}"
+
+
+# --- Gemma Story Generation Function ---
+# Replace the old generation function with one specific to Gemma-IT
+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 to the prompt.
+    """
+    if tokenizer is None or model is None:
+        raise RuntimeError("Model and tokenizer not loaded. App startup failed?")
+
+    # Gemma-IT uses a specific chat template. We wrap the user's prompt in it.
+    messages = [
+        {"role": "user", "content": prompt_text}
+        # You could add a system prompt here if desired, but Gemma-IT
+        # often works well with a detailed user prompt.
+    ]
+
+    # Apply the chat template. This adds the necessary special tokens
+    # and formatting for the model to understand the instruction.
+    # `add_generation_prompt=True` adds the token that signals the model
+    # should start generating its response.
+    input_text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False, # Keep as string for encoding later
+        add_generation_prompt=True # Add the assistant turn prompt
+    )
 
+    # Encode the templated prompt
+    # Max length should consider the prompt length + generated length
+    # Max input context for Gemma is 8192 tokens, but keeping prompt shorter is better for CPU
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=1024) # Using a reasonable max_length for input
 
-import random
+    # Ensure inputs are on the correct device (CPU by default)
+    # inputs = {k: v.to(model.device) for k, v in inputs.items()} # Redundant on CPU
 
-from fastapi import Query
-from deep_translator import GoogleTranslator
-from deep_translator.exceptions import InvalidSourceOrTargetLanguage
+    # Generate text
+    # The generate method returns the input_ids plus the generated tokens
+    generate_ids = model.generate(
+        inputs.input_ids,
+        max_new_tokens=max_new_tokens,
+        do_sample=True, # Set to True for creative text generation
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        pad_token_id=tokenizer.pad_token_id, # Use the pad token during generation
+        # Gemma's EOS token is handled by default generate logic often
+        # eos_token_id=tokenizer.eos_token_id
+    )
 
+    # Decode the generated text.
+    # We slice generate_ids to exclude the input prompt tokens, only decoding the new ones.
+    # The slicing [0, inputs.input_ids.shape[-1]:] selects the generated part for the first (and only) item in the batch
+    # The `skip_special_tokens=True` removes special tokens like <start_of_turn>, <end_of_turn>, <eos>
+    generated_text = tokenizer.decode(generate_ids[0, inputs.input_ids.shape[-1]:], skip_special_tokens=True)
 
-from fastapi import Query
+    # Gemma responses might sometimes include extra whitespace or turn markers if decoding is not perfect.
+    # Further cleanup might be needed depending on the exact output format, but skip_special_tokens helps.
+    # We can also remove leading/trailing whitespace.
+    return generated_text.strip()
 
+# --- FastAPI Endpoint ---
 @app.post("/generate-story/")
-async def generate_story(image_file: UploadFile = File(...), language: str = Form(...)):
+async def generate_story_endpoint(image_file: UploadFile = File(...), language: str = Form(...)):
     image_data = await image_file.read()
-    system_prompt = f"write an attractive story in 300 words about {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', 'a journey into the unknown'])}"
 
+    # Choose a random theme for the story prompt
+    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',
+        'a journey into the unknown'
+    ])
+
+    # Get image caption
     caption = generate_image_caption(image_data)
     if caption.startswith("Error"):
+        print(f"Caption generation failed: {caption}")
         raise HTTPException(status_code=500, detail=caption)
-    ai_response = next(run(caption, [], system_prompt))
-
-    if language != "english":
-        translator = GoogleTranslator(source='english', target=language)
-        ai_response = translator.translate(ai_response)
 
-    return {"story": ai_response}
+    # Construct the detailed prompt for Gemma-IT.
+    # Instruct it clearly to write a story based on the theme and incorporating 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 local Gemma model
+    try:
+        story = generate_story_gemma(
+            prompt_text,
+            max_new_tokens=300, # Generate up to 300 new tokens
+            temperature=0.7,    # Controls randomness (higher = more random)
+            top_p=0.9,          # Controls diversity (nucleus sampling)
+            top_k=50            # Controls diversity (top-k sampling)
+        )
+        # Basic cleanup: Sometimes models might start with whitespace or unwanted characters
+        story = story.strip()
+
+    except Exception as e:
+        print(f"Story generation failed: {e}") # Log generation errors
+        # Provide more detail in the HTTP exception for debugging
+        raise HTTPException(status_code=500, detail=f"Story generation failed: {e}. Please check Space logs for details.")
+
+
+    # Translate the story if the target language is not English
+    if language.lower() != "english":
+        try:
+            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}")
+                 # Return English story with a warning
+                 return {"story": story + "\n\n(Note: Automatic translation to your requested language failed.)"}
+
+            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:
+             print(f"Translation failed for language {language}: {e}") # Log translation errors
+             raise HTTPException(status_code=500, detail=f"Translation failed: {e}")
+
+
+    # Return the generated (and potentially translated) story
+    return {"story": story}
+
+# --- Optional: Serve a simple HTML form for testing (Needs templates dir and 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):
+#     return templates.TemplateResponse("index.html", {"request": request})