Update app.py

app.py

@@ -1,12 +1,12 @@
 import os
 import torch
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, field_validator
 from transformers import (
     AutoConfig,
     pipeline,
-    AutoModelForCausalLM,
+    AutoModelForSeq2SeqLM, # Changed AutoModelForCausalLM to AutoModelForSeq2SeqLM
     AutoTokenizer,
     GenerationConfig,
     StoppingCriteriaList
@@ -69,7 +69,7 @@ class S3ModelLoader:
         s3_uri = self._get_s3_uri(model_name)
         try:
             config = AutoConfig.from_pretrained(s3_uri, local_files_only=True)
-            model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config, local_files_only=True)
+            model = AutoModelForSeq2SeqLM.from_pretrained(s3_uri, config=config, local_files_only=True) # Changed AutoModelForCausalLM
             tokenizer = AutoTokenizer.from_pretrained(s3_uri, config=config, local_files_only=True)
 
             if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
@@ -80,7 +80,7 @@ class S3ModelLoader:
             try:
                 config = AutoConfig.from_pretrained(model_name)
                 tokenizer = AutoTokenizer.from_pretrained(model_name, config=config)
-                model = AutoModelForCausalLM.from_pretrained(model_name, config=config)
+                model = AutoModelForSeq2SeqLM.from_pretrained(model_name, config=config) # Changed AutoModelForCausalLM
 
                 if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
                     tokenizer.pad_token_id = config.pad_token_id or tokenizer.eos_token_id
@@ -135,7 +135,6 @@ async def generate(request: GenerateRequest):
         raise HTTPException(status_code=500,
                             detail=f"Internal server error: {str(e)}")
 
-
 async def stream_text(model, tokenizer, input_text,
                         generation_config, stop_sequences,
                         device, chunk_delay, max_length=2048):
@@ -199,8 +198,6 @@ async def stream_text(model, tokenizer, input_text,
                                truncation=True,
                                max_length=max_length).to(device)
 
-        
-
 @app.post("/generate-image")
 async def generate_image(request: GenerateRequest):
     try: