Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 import requests
-import json
 import io
 import random
 import os
@@ -39,13 +38,12 @@ def extend_prompt(input_text):
         response = requests.post(API_URL, headers=HEADERS, json=payload)
         response.raise_for_status()
 
-        output = response.json()
-
-        extended_prompt = json.loads(output)[0]["generated_text"].strip()
-        return extended_prompt
+        output = response.json()[0]
+        return output
         
     except requests.exceptions.RequestException as e:
-        return None
+        gr.Warning("Cannot extend prompt: Please try again.")
+        return input_text
 
 def generate_image(prompt, selected_model, is_negative=False, steps=1, cfg_scale=6, seed=None):
     if not prompt.strip():