Update app.py

app.py

@@ -121,17 +121,16 @@ async def generate(request: GenerateRequest, model_resources: tuple = Depends(ge
                 raise HTTPException(status_code=400, detail="Model mismatch for continuation.")
             input_text = previous_data["output"]
 
-        generation_config = GenerationConfig(
-            temperature=temperature,
-            max_new_tokens=max_new_tokens,
-            top_p=top_p,
-            top_k=top_k,
-            repetition_penalty=repetition_penalty,
-            do_sample=do_sample,
-            num_return_sequences=num_return_sequences,
-            no_repeat_ngram_size=no_repeat_ngram_size,
-            pad_token_id=tokenizer.pad_token_id
-        )
+        generation_config = GenerationConfig.from_pretrained(model_name) # Load default config and override
+        generation_config.temperature = temperature
+        generation_config.max_new_tokens = max_new_tokens
+        generation_config.top_p = top_p
+        generation_config.top_k = top_k
+        generation_config.repetition_penalty = repetition_penalty
+        generation_config.do_sample = do_sample
+        generation_config.num_return_sequences = num_return_sequences
+        generation_config.no_repeat_ngram_size = no_repeat_ngram_size
+        generation_config.pad_token_id = tokenizer.pad_token_id
 
         generated_text = generate_text_internal(model, tokenizer, input_text, generation_config, stop_sequences)
 
@@ -147,19 +146,23 @@ async def generate(request: GenerateRequest, model_resources: tuple = Depends(ge
 
 def generate_text_internal(model, tokenizer, input_text, generation_config, stop_sequences):
     max_model_length = model.config.max_position_embeddings
-    encoded_input = tokenizer(input_text, return_tensors="pt", max_length=max_model_length, truncation=True)
+    encoded_input = tokenizer(input_text, return_tensors="pt", max_length=max_model_length, truncation=True).to(model.device) # Ensure input is on the same device as the model
 
     stopping_criteria = StoppingCriteriaList()
 
-    class CustomStoppingCriteria(StoppingCriteriaList):
+    class CustomStoppingCriteria(StoppingCriteria): # Inherit directly from StoppingCriteria
+        def __init__(self, stop_sequences, tokenizer):
+            self.stop_sequences = stop_sequences
+            self.tokenizer = tokenizer
+
         def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-            decoded_output = tokenizer.decode(input_ids[0], skip_special_tokens=True)
-            for stop in stop_sequences:
+            decoded_output = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)
+            for stop in self.stop_sequences:
                 if decoded_output.endswith(stop):
                     return True
             return False
 
-    stopping_criteria.append(CustomStoppingCriteria())
+    stopping_criteria.append(CustomStoppingCriteria(stop_sequences, tokenizer))
 
     outputs = model.generate(
         encoded_input.input_ids,
@@ -174,7 +177,7 @@ def generate_text_internal(model, tokenizer, input_text, generation_config, stop
 async def load_pipeline_from_s3(task, model_name):
     s3_uri = f"s3://{S3_BUCKET_NAME}/{model_name.replace('/', '-')}"
     try:
-        return pipeline(task, model=s3_uri)
+        return pipeline(task, model=s3_uri, token=HUGGINGFACE_HUB_TOKEN) # Include token if needed
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error loading {task} model from S3: {e}")
 
@@ -186,9 +189,11 @@ 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]
+        image_path = f"generated_image_{os.urandom(8).hex()}.png" # Save image locally
+        image.save(image_path)
         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, "model_name": request.model_name})
+        active_generations[new_continuation_id] = {"model_name": request.model_name, "output": f"Image saved to {image_path}"} # Return path or upload URL
+        return JSONResponse({"url": image_path, "continuation_id": new_continuation_id, "model_name": request.model_name})
 
     except HTTPException as http_err:
         raise http_err
@@ -202,10 +207,12 @@ async def generate_text_to_speech(request: GenerateRequest):
             raise HTTPException(status_code=400, detail="Invalid task_type for this endpoint.")
 
         tts_pipeline = await load_pipeline_from_s3("text-to-speech", request.model_name)
-        output = tts_pipeline(request.input_text)
+        audio_output = tts_pipeline(request.input_text)
+        audio_path = f"generated_audio_{os.urandom(8).hex()}.wav"
+        sf.write(audio_path, audio_output["sampling_rate"], audio_output["audio"])
         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, "model_name": request.model_name})
+        active_generations[new_continuation_id] = {"model_name": request.model_name, "output": f"Audio saved to {audio_path}"}
+        return JSONResponse({"url": audio_path, "continuation_id": new_continuation_id, "model_name": request.model_name})
 
     except HTTPException as http_err:
         raise http_err
@@ -219,10 +226,12 @@ async def generate_video(request: GenerateRequest):
             raise HTTPException(status_code=400, detail="Invalid task_type for this endpoint.")
 
         video_pipeline = await load_pipeline_from_s3("text-to-video", request.model_name)
-        output = video_pipeline(request.input_text)
+        video_frames = video_pipeline(request.input_text).frames
+        video_path = f"generated_video_{os.urandom(8).hex()}.mp4"
+        imageio.mimsave(video_path, video_frames, fps=30) # Adjust fps as needed
         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, "model_name": request.model_name})
+        active_generations[new_continuation_id] = {"model_name": request.model_name, "output": f"Video saved to {video_path}"}
+        return JSONResponse({"url": video_path, "continuation_id": new_continuation_id, "model_name": request.model_name})
 
     except HTTPException as http_err:
         raise http_err