Update app.py

app.py

@@ -4,6 +4,8 @@ import json
 import io
 import random
 import os
+import torch
+from transformers import pipeline
 from PIL import Image
 
 API_BASE_URL = "https://api-inference.huggingface.co/models/"
@@ -28,6 +30,8 @@ MODEL_LIST = [
 API_TOKEN = os.getenv("HF_READ_TOKEN")  # Make sure to set your Hugging Face token
 HEADERS = {"Authorization": f"Bearer {API_TOKEN}"}
 
+pipe = pipeline("text-generation", model="isek-ai/SDPrompt-RetNet-300M", trust_remote_code=True)
+
 def select_model(model_name):
     if model_name in MODEL_LIST:
         return f"{API_BASE_URL}{model_name}"
@@ -37,17 +41,12 @@ def extend_prompt(input_text):
         gr.Warning("Input text is empty!")
         return None
         
-    API_URL = f"{API_BASE_URL}Gustavosta/MagicPrompt-Stable-Diffusion"
+    seed = random.randint(1, 1000000)
+    torch.manual_seed(seed)
     
-    payload = {"inputs": input_text}
+    output = pipe(input_text, max_length=(len(input_text) + random.randint(60, 90)), num_return_sequences=4)
 
-    try:
-        response = requests.post(API_URL, headers=HEADERS, json=payload)
-        response.raise_for_status()
-        return response.json()[0].get("generated_text", "")
-    except requests.exceptions.RequestException as e:
-        gr.Warning(f"Error in API request: {e}")
-        return None
+    return output
 
 def generate_image(prompt, selected_model, is_negative=False, steps=1, cfg_scale=6, seed=None):
     if not prompt.strip():