Update app.py

app.py

@@ -1,52 +1,47 @@
 import gradio as gr
-from diffusers import StableDiffusionImg2ImgPipeline
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
-
-# --- Device selection: CUDA if available, else CPU ---
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# --- Load Stable Diffusion for style transfer ---
-sd_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5"
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+
+# Available models
+MODEL_OPTIONS = {
+    "Prithivida GEC v1": "prithivida/grammar_error_correcter_v1",
+    "Hassaanik GEC": "hassaanik/grammar-correction-model",
+    "Vennify T5 GEC": "vennify/t5-base-grammar-correction"
+}
+
+# Cache loaded pipelines so we don’t reload every time
+loaded_pipelines = {}
+
+def get_pipeline(model_id):
+    if model_id not in loaded_pipelines:
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+        loaded_pipelines[model_id] = pipeline("text2text-generation",
+                                              model=model,
+                                              tokenizer=tokenizer)
+    return loaded_pipelines[model_id]
+
+def oxford_polish(sentence: str, model_choice: str) -> str:
+    model_id = MODEL_OPTIONS[model_choice]
+    polisher = get_pipeline(model_id)
+    prompt = (
+        "Correct this sentence into formal written English, following the Oxford University Style Guide. "
+        "Ensure tense matches time expressions (e.g. 'tomorrow' → future, 'yesterday' → past), "
+        "use British spelling, apply the Oxford comma, and correct uncountable nouns naturally. "
+        "Sentence: " + sentence
+    )
+    out = polisher(prompt, max_new_tokens=80, do_sample=False)
+    return out[0]["generated_text"].strip()
+
+# Gradio interface
+demo = gr.Interface(
+    fn=oxford_polish,
+    inputs=[
+        gr.Textbox(lines=2, placeholder="Enter a sentence to correct..."),
+        gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), value="Prithivida GEC v1", label="Choose Model")
+    ],
+    outputs=gr.Textbox(label="Oxford-style Correction"),
+    title="Oxford Grammar Polisher",
+    description="Test multiple free grammar correction models from Hugging Face Hub with Oxford-style rules."
 )
-sd_pipe = sd_pipe.to(device)
-sd_pipe.enable_attention_slicing()  # helps reduce memory usage
-
-# --- Load T5-base or T5-small for grammar correction ---
-model_name = "vennify/t5-base-grammar-correction"  # or "prithivida/grammar_error_correcter_v1"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
-
-def style_transfer(input_image, prompt, strength=0.5, guidance=7.5):
-    result = sd_pipe(
-        prompt=prompt,
-        image=input_image,
-        strength=strength,
-        guidance_scale=guidance
-    ).images[0]
-    return result
-
-def correct_text(text):
-    inputs = tokenizer("gec: " + text, return_tensors="pt").to(device)
-    outputs = model.generate(**inputs, max_length=128)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# --- Gradio UI ---
-with gr.Blocks() as demo:
-    gr.Markdown("## 🎨 Style Transfer + ✍️ Grammar Correction")
-
-    with gr.Tab("Image Style Transfer"):
-        img_in = gr.Image(type="pil")
-        prompt = gr.Textbox(label="Style Prompt")
-        img_out = gr.Image()
-        btn1 = gr.Button("Transfer Style")
-        btn1.click(style_transfer, inputs=[img_in, prompt], outputs=img_out)
-
-    with gr.Tab("Text Correction"):
-        txt_in = gr.Textbox(label="Enter text")
-        txt_out = gr.Textbox(label="Corrected text")
-        btn2 = gr.Button("Correct")
-        btn2.click(correct_text, inputs=txt_in, outputs=txt_out)
 
 demo.launch()