Update app.py

app.py

@@ -1,74 +1,105 @@
 import gradio as gr
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import os
+import subprocess
+import zipfile
+import requests
+from huggingface_hub import hf_hub_download
 
-MODEL_ID = "tencent/HY-MT1.5-1.8B"
+BASE_DIR = os.getcwd()
+SDCPP_DIR = os.path.join(BASE_DIR, "sdcpp")
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.bfloat16 if device == "cuda" else torch.float32
+# =========================
+# 1. دانلود stable-diffusion.cpp
+# =========================
+def setup_sdcpp():
+    if not os.path.exists(SDCPP_DIR):
+        os.makedirs(SDCPP_DIR, exist_ok=True)
 
-# load model + tokenizer
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+        zip_url = "https://github.com/leejet/stable-diffusion.cpp/releases/download/master-586-c97702e/sd-master-c97702e-bin-Linux-Ubuntu-24.04-x86_64.zip"
+        zip_path = os.path.join(BASE_DIR, "sdcpp.zip")
 
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=dtype,
-    device_map="cpu",
-    trust_remote_code=True
-)
+        print("Downloading stable-diffusion.cpp...")
+        r = requests.get(zip_url)
+        with open(zip_path, "wb") as f:
+            f.write(r.content)
 
-model.eval()
+        print("Extracting...")
+        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+            zip_ref.extractall(SDCPP_DIR)
 
-# inference function
-def generate(text, target_language):
-    if not text.strip():
-        return ""
+        os.remove(zip_path)
 
-    messages = [
-        {"role": "user", "content": f"Translate the following segment into {target_language}, without additional explanation.\n\n{text}"},
-    ]
-    
-    tokenized_chat = tokenizer.apply_chat_template(
-        messages,
-        tokenize=True,
-        add_generation_prompt=True,
-        return_tensors="pt"
+        # chmod
+        subprocess.run(["chmod", "+x", f"{SDCPP_DIR}/sd-cli"])
+
+
+# =========================
+# 2. دانلود مدل‌ها
+# =========================
+def setup_models():
+    print("Downloading models...")
+
+    model_path = hf_hub_download(
+        repo_id="leejet/Z-Image-Turbo-GGUF",
+        filename="z-image-turbo-Q4_K_M.gguf"
     )
-    
-    outputs = model.generate(**tokenized_chat, max_new_tokens=2048)
-
-    input_length = tokenized_chat["input_ids"].shape[1]
-    generated_tokens = outputs[0][input_length:]
-    result = tokenizer.decode(generated_tokens, skip_special_tokens=True)
-
-    return result
-
-LANG_MAP = {
-    "English": "English",
-    "Persian (Farsi)": "Persian",
-    "Arabic": "Arabic",
-    "French": "French",
-    "German": "German",
-    "Spanish": "Spanish",
-    "Chinese": "Chinese",
-    "Japanese": "Japanese",
-}
 
+    vae_path = hf_hub_download(
+        repo_id="black-forest-labs/FLUX.1-schnell",
+        filename="ae.safetensors"
+    )
+
+    llm_path = hf_hub_download(
+        repo_id="unsloth/Qwen3-4B-Instruct-2507-GGUF",
+        filename="Qwen3-4B-Instruct-Q4_K_M.gguf"
+    )
+
+    return model_path, vae_path, llm_path
+
+
+MODEL_PATH, VAE_PATH, LLM_PATH = None, None, None
+
+# =========================
+# 3. inference
+# =========================
+def generate(prompt):
+    output_path = os.path.join(BASE_DIR, "output.png")
+
+    cmd = [
+        f"{SDCPP_DIR}/sd-cli",
+        "-m", MODEL_PATH,
+        "--vae", VAE_PATH,
+        "--llm", LLM_PATH,
+        "-p", prompt,
+        "--steps", "6",
+        "--cfg-scale", "1.0",
+        "-o", output_path
+    ]
+
+    env = os.environ.copy()
+    env["LD_LIBRARY_PATH"] = SDCPP_DIR
+
+    print("Running:", " ".join(cmd))
+
+    subprocess.run(cmd, env=env, check=True)
+
+    return output_path
+
+
+# =========================
+# init
+# =========================
+setup_sdcpp()
+MODEL_PATH, VAE_PATH, LLM_PATH = setup_models()
+
+# =========================
 # UI
+# =========================
 demo = gr.Interface(
     fn=generate,
-    inputs=[
-        gr.Textbox(label="Source text", lines=5),
-
-        gr.Dropdown(
-            choices=list(LANG_MAP.keys()),
-            value="English",
-            label="Target Language 🌍"
-        ),
-    ],
-    outputs=gr.Textbox(label="Output"),
-    title="HY-MT1.5-1.8B Demo",
-    description="Multilingual text generation test"
+    inputs=gr.Textbox(label="Prompt"),
+    outputs=gr.Image(label="Generated Image"),
+    title="Z-Image Turbo GGUF (CPU)"
 )
 
 demo.launch()