Update app.py

app.py

@@ -1,42 +1,121 @@
 import gradio as gr
-import spaces
+import gc
+import os
+import random
+import tempfile
+from pathlib import Path
+
+import numpy as np
 import torch
-from threading import Thread
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from PIL import Image
 from fastapi.responses import HTMLResponse
-from pathlib import Path
+from gradio.data_classes import FileData
+
+# ZeroGPU. Degrade gracefully off-Spaces so `python app.py` works locally.
+try:
+    import spaces
+    _HAS_SPACES = True
+except ImportError:
+    _HAS_SPACES = False
+
+# --- Model load ---------------------------------------------------------------
+# Heavy startup is wrapped in `gr.NO_RELOAD` so `gradio app.py` hot reload
+# does not redownload weights every time you save the HTML.
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.bfloat16
+
+if gr.NO_RELOAD:
+    from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
+    from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
+    from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
+
+    PIPE = QwenImageEditPlusPipeline.from_pretrained(
+        "FireRedTeam/FireRed-Image-Edit-1.1",
+        transformer=QwenImageTransformer2DModel.from_pretrained(
+            "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19",
+            torch_dtype=DTYPE,
+            device_map="cuda",
+        ),
+        torch_dtype=DTYPE,
+    ).to(DEVICE)
+
+    try:
+        PIPE.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+        print("Flash Attention 3 processor set.")
+    except Exception as e:
+        print(f"FA3 processor not set: {e}")
+
+NEGATIVE_PROMPT = (
+    "worst quality, low quality, bad anatomy, bad hands, text, error, "
+    "missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, "
+    "signature, watermark, username, blurry"
+)
+MAX_SEED = np.iinfo(np.int32).max
+
+
+def _round_dims(image: Image.Image) -> tuple[int, int]:
+    w, h = image.size
+    if w > h:
+        new_w, new_h = 1024, int(1024 * h / w)
+    else:
+        new_h, new_w = 1024, int(1024 * w / h)
+    return (new_w // 8) * 8, (new_h // 8) * 8
 
-app = gr.Server()
+
+# --- Inner GPU function -------------------------------------------------------
+# Per the reference: @spaces.GPU goes on the *inner* function that runs the
+# model. The outer @server.api route just plugs it into the queue.
+if _HAS_SPACES:
+    @spaces.GPU
+    def _edit(image: Image.Image, prompt: str, seed: int, steps: int) -> Image.Image:
+        return _run_pipe(image, prompt, seed, steps)
+else:
+    def _edit(image, prompt, seed, steps):
+        return _run_pipe(image, prompt, seed, steps)
+
+
+def _run_pipe(image, prompt, seed, steps):
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    width, height = _round_dims(image)
+    generator = torch.Generator(device=DEVICE).manual_seed(seed)
+    return PIPE(
+        image=[image],
+        prompt=prompt,
+        negative_prompt=NEGATIVE_PROMPT,
+        width=width,
+        height=height,
+        num_inference_steps=steps,
+        true_cfg_scale=1.0,
+        generator=generator,
+    ).images[0]
+
+
+# --- Server -------------------------------------------------------------------
+server = gr.Server()
 HOME = Path(__file__).parent
 
-MODEL_ID = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
-tok = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16).to("cuda")
-
-@spaces.GPU
-def _generate(text: str):
-    inputs = tok.apply_chat_template(
-        [{"role": "user", "content": f"Summarize in 3 bullets:\n\n{text}"}],
-        return_tensors="pt", return_dict=True, add_generation_prompt=True,
-    ).to("cuda")
-    streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True)
-    Thread(target=model.generate, kwargs=dict(
-        **inputs, streamer=streamer, max_new_tokens=300, do_sample=False,
-    )).start()
-    return streamer
-
-@app.mcp.tool(name="summarize")
-@app.api(name="summarize", concurrency_limit=1, stream_every=0.2)
-def summarize(text: str) -> str:
-    """Summarize the input text into 3 bullet points."""
-    out = ""
-    for chunk in _generate(text):
-        out += chunk
-        yield out
-
-@app.get("/", response_class=HTMLResponse)
-async def index():
+
+@server.api(name="edit_image", concurrency_limit=1)
+def edit_image(image: FileData, prompt: str) -> dict:
+    """Edit an image guided by a text prompt using FireRed-Image-Edit 1.1."""
+    if not prompt or not prompt.strip():
+        return {"error": "Please enter an edit prompt."}
+    src = Image.open(image["path"]).convert("RGB")
+    seed = random.randint(0, MAX_SEED)
+    result = _edit(src, prompt.strip(), seed, steps=4)
+
+    fd, out_path = tempfile.mkstemp(suffix=".png")
+    os.close(fd)
+    result.save(out_path)
+    return {"image": FileData(path=out_path), "seed": seed}
+
+
+@server.get("/", response_class=HTMLResponse)
+async def homepage():
     return (HOME / "index.html").read_text(encoding="utf-8")
 
+
 if __name__ == "__main__":
-    app.launch(mcp_server=True)
+    server.launch(mcp_server=True, show_error=True)