Update app.py

app.py

@@ -1,52 +1,28 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-import torch
 import gradio as gr
-
-# --- Load model in 8-bit mode on CPU ---
-print("Loading model in 8-bit on CPU (reduced memory)...")
-
-quant_config = BitsAndBytesConfig(
-    load_in_8bit=True,
-    llm_int8_enable_fp32_cpu_offload=True,  # partial CPU offload
-)
-
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(
-    "m-a-p/YuE-s1-7B-anneal-en-cot",
-    trust_remote_code=True
-)
-
-# Load model
-model = AutoModelForCausalLM.from_pretrained(
-    "m-a-p/YuE-s1-7B-anneal-en-cot",
-    quantization_config=quant_config,
-    device_map="cpu",        # force CPU usage
-    torch_dtype=torch.float16,
-    attn_implementation="eager",
-    trust_remote_code=True
-)
-
-model.eval()
-print("✅ Model loaded (8-bit, ~6–8 GB RAM).")
-
-# --- Define generation function ---
-def generate_text(prompt, max_tokens=200):
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+MODEL_NAME = "m-a-p/YuE-s1-7B-anneal-en-cot"
+
+def generate_text(prompt):
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float32,
+        attn_implementation="eager"
+    )
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    
     inputs = tokenizer(prompt, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model.generate(**inputs, max_new_tokens=int(max_tokens))
+    outputs = model.generate(**inputs, max_new_tokens=200)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# --- Gradio Interface ---
 demo = gr.Interface(
     fn=generate_text,
-    inputs=[
-        gr.Textbox(label="Enter your prompt", lines=4, placeholder="Type something..."),
-        gr.Slider(10, 400, value=200, step=10, label="Max tokens")
-    ],
-    outputs=gr.Textbox(label="Generated Text"),
-    title="YuE-s1-7B Text Generator (8-bit CPU)",
-    description="A lightweight quantized version running on CPU (~6–8 GB RAM)."
+    inputs=gr.Textbox(lines=3, label="Input Prompt"),
+    outputs="text",
+    title="YuE Text Generator",
+    description="A Hugging Face Space for generating text using YuE 7B model."
 )
 
 if __name__ == "__main__":
-    demo.launch(server_name="127.0.0.1", server_port=7860)
+    demo.launch()