Update app.py

app.py

@@ -6,7 +6,7 @@ from fastapi.middleware.cors import CORSMiddleware
 import io
 import torch
 from diffusers import AudioLDM2Pipeline
-from transformers import GPT2LMHeadModel  # <-- IMPORT THE CORRECT MODEL TYPE
+from transformers import GPT2LMHeadModel
 from scipy.io.wavfile import write as write_wav
 import numpy as np
 
@@ -35,26 +35,28 @@ try:
 
     # --- FIX for Model Component Mismatch ---
     # 1. Manually load the correct GPT2 model variant (GPT2LMHeadModel).
-    # The sub-model 'gpt2' is used by audioldm2 for prompt understanding.
     logger.info("Pre-loading the correct language model component (GPT2LMHeadModel)...")
     language_model = GPT2LMHeadModel.from_pretrained(
         "openai-community/gpt2", cache_dir=cache_dir
     ).to(device)
     logger.info("Language model component loaded successfully.")
 
+    # ########################################################################## #
+    # ############## THE FINAL FIX: Add the missing configuration ############## #
+    language_model.config.max_new_tokens = 32 # This line is the fix.
+    # ########################################################################## #
+
     # 2. Load the main pipeline, injecting our pre-loaded component.
-    # This forces the pipeline to use the correct model and avoid the AttributeError.
     pipeline = AudioLDM2Pipeline.from_pretrained(
         MODEL_REPO,
         torch_dtype=torch_dtype,
         cache_dir=cache_dir,
-        language_model=language_model,  # <-- INJECT THE CORRECT COMPONENT
+        language_model=language_model,
     )
     pipeline = pipeline.to(device)
     logger.info("Model loaded successfully and moved to device.")
 except Exception as e:
     logger.error(f"Fatal error during model loading: {e}", exc_info=True)
-    # If the model fails to load, the 'pipeline' variable will remain None.
 
 # --- CORS Middleware ---
 app.add_middleware(