Update app.py

app.py

@@ -7,6 +7,7 @@ from transformers import (
     AutoConfig,
     pipeline,
     AutoModelForSeq2SeqLM,
+    AutoModelForCausalLM,
     AutoTokenizer,
     GenerationConfig,
     StoppingCriteriaList
@@ -14,10 +15,11 @@ from transformers import (
 import boto3
 import uvicorn
 import asyncio
-from io import BytesIO
 from transformers import pipeline
 import json
 from huggingface_hub import login
+import base64
+
 
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
@@ -80,9 +82,20 @@ class S3ModelLoader:
             config = AutoConfig.from_pretrained(
                 s3_uri, local_files_only=True
             )
-            model = AutoModelForSeq2SeqLM.from_pretrained(
+            
+            if "llama" in model_name:
+                model = AutoModelForCausalLM.from_pretrained(
+                s3_uri, config=config, local_files_only=True, rope_scaling = {"type": "linear", "factor": 2.0}
+            )
+            elif 'qwen' in model_name:
+                model = AutoModelForCausalLM.from_pretrained(
+                s3_uri, config=config, local_files_only=True
+            )
+            else:
+                model = AutoModelForSeq2SeqLM.from_pretrained(
                 s3_uri, config=config, local_files_only=True
             )
+                
             tokenizer = AutoTokenizer.from_pretrained(
                 s3_uri, config=config, local_files_only=True
             )
@@ -101,7 +114,17 @@ class S3ModelLoader:
                 tokenizer = AutoTokenizer.from_pretrained(
                     model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
                 )
-                model = AutoModelForSeq2SeqLM.from_pretrained(
+                
+                if "llama" in model_name:
+                    model = AutoModelForCausalLM.from_pretrained(
+                    model_name, config=config, token=HUGGINGFACE_HUB_TOKEN, rope_scaling = {"type": "linear", "factor": 2.0}
+                )
+                elif 'qwen' in model_name:
+                    model = AutoModelForCausalLM.from_pretrained(
+                    model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
+                )
+                else:
+                    model = AutoModelForSeq2SeqLM.from_pretrained(
                     model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
                 )
 
@@ -261,14 +284,10 @@ async def generate_image(request: GenerateRequest):
             device=device
         )
         image = image_generator(validated_body.input_text)[0]
-
-        img_byte_arr = BytesIO()
-        image.save(img_byte_arr, format="PNG")
-        img_byte_arr.seek(0)
-
-        return StreamingResponse(
-            img_byte_arr, media_type="image/png"
-        )
+        
+        image_data = list(image.getdata())
+        
+        return json.dumps({"image_data": image_data, "is_end": True})
 
     except Exception as e:
         raise HTTPException(
@@ -287,15 +306,14 @@ async def generate_text_to_speech(request: GenerateRequest):
             "text-to-speech", model=validated_body.model_name,
             device=device
         )
-        audio = audio_generator(validated_body.input_text)[0]
+        audio = audio_generator(validated_body.input_text)
 
-        audio_byte_arr = BytesIO()
-        audio.save(audio_byte_arr)
-        audio_byte_arr.seek(0)
-
-        return StreamingResponse(
-            audio_byte_arr, media_type="audio/wav"
-        )
+        
+        audio_bytes = audio["audio"]
+        
+        audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
+        
+        return json.dumps({"audio": audio_base64, "is_end": True})
 
     except Exception as e:
         raise HTTPException(
@@ -313,15 +331,12 @@ async def generate_video(request: GenerateRequest):
             "text-to-video", model=validated_body.model_name,
             device=device
         )
-        video = video_generator(validated_body.input_text)[0]
-
-        video_byte_arr = BytesIO()
-        video.save(video_byte_arr)
-        video_byte_arr.seek(0)
-
-        return StreamingResponse(
-            video_byte_arr, media_type="video/mp4"
-        )
+        video = video_generator(validated_body.input_text)
+        
+        
+        video_base64 = base64.b64encode(video).decode('utf-8')
+        
+        return json.dumps({"video": video_base64, "is_end": True})
 
     except Exception as e:
         raise HTTPException(