Update app.py

app.py

@@ -7,7 +7,7 @@ import boto3
 import uvicorn
 import soundfile as sf
 import imageio
-from typing import Dict
+from typing import Dict, Optional
 
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
@@ -42,7 +42,7 @@ class GenerateRequest(BaseModel):
     do_sample: bool = True
     stop_sequences: list[str] = []
     no_repeat_ngram_size: int = 2
-    continuation_id: str = None
+    continuation_id: Optional[str] = None
 
     @field_validator("model_name")
     def model_name_cannot_be_empty(cls, v):
@@ -115,8 +115,10 @@ async def generate(request: GenerateRequest, model_resources: tuple = Depends(ge
         if continuation_id:
             if continuation_id not in active_generations:
                 raise HTTPException(status_code=404, detail="Continuation ID not found.")
-            previous_output = active_generations[continuation_id]["output"]
-            input_text = previous_output
+            previous_data = active_generations[continuation_id]
+            if previous_data["model_name"] != model_name:
+                raise HTTPException(status_code=400, detail="Model mismatch for continuation.")
+            input_text = previous_data["output"]
 
         generation_config = GenerationConfig(
             temperature=temperature,
@@ -132,13 +134,10 @@ async def generate(request: GenerateRequest, model_resources: tuple = Depends(ge
 
         generated_text = generate_text_internal(model, tokenizer, input_text, generation_config, stop_sequences)
 
-        if not continuation_id:
-            continuation_id = os.urandom(16).hex()
-            active_generations[continuation_id] = {"model_name": model_name, "output": generated_text}
-        else:
-            active_generations[continuation_id]["output"] = generated_text
+        new_continuation_id = continuation_id if continuation_id else os.urandom(16).hex()
+        active_generations[new_continuation_id] = {"model_name": model_name, "output": generated_text}
 
-        return JSONResponse({"text": generated_text, "continuation_id": continuation_id})
+        return JSONResponse({"text": generated_text, "continuation_id": new_continuation_id})
 
     except HTTPException as http_err:
         raise http_err
@@ -186,9 +185,9 @@ async def generate_image(request: GenerateRequest):
 
         image_generator = await load_pipeline_from_s3("text-to-image", request.model_name)
         image = image_generator(request.input_text)[0]
-        continuation_id = os.urandom(16).hex()
-        active_generations[continuation_id] = {"model_name": request.model_name, "output": "Image generated successfully"}
-        return JSONResponse({"url": "Image generated successfully", "continuation_id": continuation_id})
+        new_continuation_id = os.urandom(16).hex()
+        active_generations[new_continuation_id] = {"model_name": request.model_name, "output": "Image generated successfully"}
+        return JSONResponse({"url": "Image generated successfully", "continuation_id": new_continuation_id})
 
     except HTTPException as http_err:
         raise http_err
@@ -203,9 +202,9 @@ async def generate_text_to_speech(request: GenerateRequest):
 
         tts_pipeline = await load_pipeline_from_s3("text-to-speech", request.model_name)
         output = tts_pipeline(request.input_text)
-        continuation_id = os.urandom(16).hex()
-        active_generations[continuation_id] = {"model_name": request.model_name, "output": "Audio generated successfully"}
-        return JSONResponse({"url": "Audio generated successfully", "continuation_id": continuation_id})
+        new_continuation_id = os.urandom(16).hex()
+        active_generations[new_continuation_id] = {"model_name": request.model_name, "output": "Audio generated successfully"}
+        return JSONResponse({"url": "Audio generated successfully", "continuation_id": new_continuation_id})
 
     except HTTPException as http_err:
         raise http_err
@@ -220,9 +219,9 @@ async def generate_video(request: GenerateRequest):
 
         video_pipeline = await load_pipeline_from_s3("text-to-video", request.model_name)
         output = video_pipeline(request.input_text)
-        continuation_id = os.urandom(16).hex()
-        active_generations[continuation_id] = {"model_name": request.model_name, "output": "Video generated successfully"}
-        return JSONResponse({"url": "Video generated successfully", "continuation_id": continuation_id})
+        new_continuation_id = os.urandom(16).hex()
+        active_generations[new_continuation_id] = {"model_name": request.model_name, "output": "Video generated successfully"}
+        return JSONResponse({"url": "Video generated successfully", "continuation_id": new_continuation_id})
 
     except HTTPException as http_err:
         raise http_err