Update app.py

app.py

@@ -1,47 +1,266 @@
+import os
+import gc
+import time
 import gradio as gr
-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)
+import torch
+from PIL import Image
+
+# -----------------------
+# Device + CPU perf knobs
+# -----------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Threads (tune for HF CPU Space)
+os.environ.setdefault("OMP_NUM_THREADS", "4")
+os.environ.setdefault("MKL_NUM_THREADS", "4")
+torch.set_num_threads(int(os.environ["OMP_NUM_THREADS"]))
+torch.set_num_interop_threads(max(1, int(int(os.environ["OMP_NUM_THREADS"]) // 2)))
+
+INFER = torch.inference_mode if hasattr(torch, "inference_mode") else torch.no_grad
+
+# -----------------------
+# Stable Diffusion 1.5 (img2img) for style transfer
+# -----------------------
+from diffusers import StableDiffusionImg2ImgPipeline, EulerAncestralDiscreteScheduler
+
+def load_sd15_pipe():
+    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        safety_checker=None,
+        requires_safety_checker=False,
+    )
+    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    pipe = pipe.to(device)
+    pipe.enable_attention_slicing()
+    pipe.enable_vae_tiling()
+    pipe.enable_vae_slicing()
+    if device == "cuda":
+        pipe.unet.to(memory_format=torch.channels_last)
+    return pipe
+
+_sd_pipe = None
+
+def sd_style_transfer(input_image, prompt, strength=0.55, guidance=5.5, steps=18, width=512, height=512, seed=0):
+    global _sd_pipe
+    if input_image is None:
+        raise gr.Error("Please upload an input image.")
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please provide a style prompt.")
+
+    if _sd_pipe is None:
+        t0 = time.time()
+        _sd_pipe = load_sd15_pipe()
+        print(f"[SD] Pipeline loaded in {time.time()-t0:.2f}s on {device}.", flush=True)
+
+    generator = torch.Generator(device=device) if device == "cuda" else torch.Generator()
+    if isinstance(seed, (int, float)) and int(seed) > 0:
+        generator = generator.manual_seed(int(seed))
+
+    img = input_image.convert("RGB").resize((int(width), int(height)), Image.LANCZOS)
+
+    with INFER():
+        out = _sd_pipe(
+            prompt=str(prompt),
+            image=img,
+            strength=float(strength),
+            guidance_scale=float(guidance),
+            num_inference_steps=int(steps),
+            generator=generator,
+        ).images[0]
+
+    if device == "cuda":
+        torch.cuda.empty_cache()
+    gc.collect()
+    return out
+
+# -----------------------
+# Grammar correction models
+# T5-small (prithivida), T5-base (vennify), GECToR (optional), Llama-3.1-8B-GEC (GGUF)
+# -----------------------
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+T5_SMALL = "prithivida/grammar_error_correcter_v1"      # T5-small
+T5_BASE  = "vennify/t5-base-grammar-correction"         # T5-base
+
+_t5_tok = {}
+_t5_mdl = {}
+
+def load_t5(model_name: str):
+    if model_name not in _t5_mdl:
+        tok = AutoTokenizer.from_pretrained(model_name)
+        mdl = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
+        _t5_tok[model_name] = tok
+        _t5_mdl[model_name] = mdl
+    return _t5_tok[model_name], _t5_mdl[model_name]
+
+def t5_correct(text: str, model_name: str, max_new_tokens=128):
+    tok, mdl = load_t5(model_name)
+    prefix = "gec: " if "prithivida" in model_name else "grammar: "
+    inputs = tok(prefix + text, return_tensors="pt").to(device)
+    with INFER():
+        out = mdl.generate(**inputs, max_length=max_new_tokens)
+    return tok.decode(out[0], skip_special_tokens=True)
+
+# ---- Optional: GECToR (lazy load) ----
+_gector_predictor = None
+_gector_error = None
+_gector_tried = False
+
+def try_load_gector():
+    global _gector_predictor, _gector_error, _gector_tried
+    if _gector_tried:
+        return _gector_predictor, _gector_error
+    _gector_tried = True
+    try:
+        from gector.gec_model import GECModel  # requires allennlp + pretrained artifacts
+        model_paths = os.environ.get("GEC_MODEL_PATHS", "").strip()
+        vocab_path = os.environ.get("GEC_VOCAB_PATH", "").strip()
+        if not model_paths or not vocab_path:
+            raise RuntimeError(
+                "GECToR selected but model artifacts are not configured. "
+                "Set GEC_MODEL_PATHS (space-separated .th files) and GEC_VOCAB_PATH (vocab dir)."
+            )
+        taggers = model_paths.split()
+        _gector_predictor = GECModel(
+            model_paths=taggers,
+            vocab_path=vocab_path,
+            device=("cuda" if device == "cuda" else "cpu"),
+            min_error_probability=0.0,
+            confidence=0.0,
+            iterations=2,
+            special_tokens_fix=1,
+        )
+    except Exception as e:
+        _gector_error = str(e)
+        _gector_predictor = None
+    return _gector_predictor, _gector_error
+
+def gector_correct(text: str):
+    predictor, err = try_load_gector()
+    if err or predictor is None:
+        return f"[GECToR not active] {err or 'Unknown error.'}\n" \
+               f"Enable by setting GEC_MODEL_PATHS and GEC_VOCAB_PATH to pretrained files."
+    tokens = text.strip().split()
+    corrected = predictor.handle_batch([tokens])[0]
+    return " ".join(corrected)
+
+# ---- Llama-3.1-8B GEC (GGUF via llama-cpp-python) ----
+_llama_model = None
+_llama_err = None
+_llama_tried = False
+
+# Choose a sensible quant filename; adjust if you upload a different one to your Space.
+LLAMA_REPO = "mradermacher/Llama-3.1-8B-Instruct-Grammatical-Error-Correction-2-GGUF"
+LLAMA_FILE = os.environ.get("LLAMA_GGUF_FILE", "llama-3.1-8b-instruct-gec.Q4_K_S.gguf")
+
+def try_load_llama():
+    global _llama_model, _llama_err, _llama_tried
+    if _llama_tried:
+        return _llama_model, _llama_err
+    _llama_tried = True
+    try:
+        from llama_cpp import Llama
+        # Load directly from Hub (no need to manually download)
+        _llama_model = Llama.from_pretrained(
+            repo_id=LLAMA_REPO,
+            filename=LLAMA_FILE,
+            n_ctx=2048,
+            n_threads=int(os.environ.get("OMP_NUM_THREADS", "4")),
+            n_batch=128,
+            verbose=False
+        )
+    except Exception as e:
+        _llama_model = None
+        _llama_err = str(e)
+    return _llama_model, _llama_err
+
+def llama_gec_correct(text: str, max_new_tokens=256):
+    mdl, err = try_load_llama()
+    if err or mdl is None:
+        return f"[Llama GGUF not active] {err or 'Unknown error.'}\n" \
+               f"Check model availability or set LLAMA_GGUF_FILE to a valid filename."
     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
+        "You are a precise grammatical error corrector. "
+        "Return only the corrected text without explanations.\n\n"
+        f"Input: {text}\n"
+        "Corrected:"
     )
-    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-grammar rules."
-)
-
-demo.launch()
+    out = mdl(prompt, max_tokens=max_new_tokens, stop=["\n\n", "\nCorrected:"])
+    return out["choices"][0]["text"].strip()
+
+# -----------------------
+# Router
+# -----------------------
+MODEL_OPTIONS = [
+    "T5-small (prithivida)",
+    "T5-base (vennify)",
+    "GECToR (tagging)",
+    "Llama-3.1-8B-GEC (GGUF)"
+]
+
+def correct_text_router(text: str, model_choice: str, max_new_tokens=128):
+    text = (text or "").strip()
+    if not text:
+        raise gr.Error("Please enter text to correct.")
+    if model_choice == "T5-small (prithivida)":
+        return t5_correct(text, T5_SMALL, max_new_tokens=max_new_tokens)
+    if model_choice == "T5-base (vennify)":
+        return t5_correct(text, T5_BASE, max_new_tokens=max_new_tokens)
+    if model_choice == "GECToR (tagging)":
+        return gector_correct(text)
+    if model_choice == "Llama-3.1-8B-GEC (GGUF)":
+        return llama_gec_correct(text, max_new_tokens=max_new_tokens)
+    return "Unknown model selection."
+
+# -----------------------
+# UI
+# -----------------------
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        f"# 🎨 Style transfer (SD 1.5 img2img) + ✍️ English correction\n"
+        f"- Device detected: **{device.upper()}**\n"
+        f"- Models: T5-small, T5-base, GECToR, Llama-3.1-8B-GEC (GGUF)\n"
+    )
+
+    with gr.Tab("Image style transfer"):
+        with gr.Row():
+            img_in = gr.Image(label="Input image", type="pil")
+            img_out = gr.Image(label="Styled output")
+        prompt = gr.Textbox(label="Style prompt", placeholder="e.g., watercolor wash, halftone dots, 1960s comic shading")
+        with gr.Row():
+            strength = gr.Slider(0.1, 0.95, value=0.55, step=0.05, label="Style strength")
+            guidance = gr.Slider(1.0, 12.0, value=5.5, step=0.5, label="Guidance")
+            steps = gr.Slider(5, 40, value=18, step=1, label="Steps")
+        with gr.Row():
+            width = gr.Slider(256, 768, value=512, step=64, label="Width")
+            height = gr.Slider(256, 768, value=512, step=64, label="Height")
+            seed = gr.Number(value=0, precision=0, label="Seed (0 = random)")
+        run_btn = gr.Button("Transfer style", variant="primary")
+        run_btn.click(
+            fn=sd_style_transfer,
+            inputs=[img_in, prompt, strength, guidance, steps, width, height, seed],
+            outputs=[img_out]
+        )
+
+    with gr.Tab("English grammar correction"):
+        model_choice = gr.Dropdown(MODEL_OPTIONS, value="T5-small (prithivida)", label="Model")
+        txt_in = gr.Textbox(lines=6, label="Input text")
+        max_new = gr.Slider(32, 512, value=128, step=16, label="Max tokens (generation models)")
+        txt_out = gr.Textbox(lines=6, label="Corrected text")
+        corr_btn = gr.Button("Correct", variant="primary")
+        corr_btn.click(
+            fn=correct_text_router,
+            inputs=[txt_in, model_choice, max_new],
+            outputs=[txt_out]
+        )
+
+    gr.Markdown(
+        "Tips:\n"
+        "- On CPU: steps 12–20, guidance 4–7, 512×512 for SD speed.\n"
+        "- T5-small = fastest, T5-base = more accurate.\n"
+        "- GECToR needs AllenNLP and pretrained tagger files (set GEC_MODEL_PATHS & GEC_VOCAB_PATH).\n"
+        "- Llama GGUF loads from Hub (Q4_K_S by default). Adjust LLAMA_GGUF_FILE if needed."
+    )
+
+if __name__ == "__main__":
+    demo.launch()