Update app.py

app.py

@@ -1,119 +1,108 @@
 import gradio as gr
-import torch
-import soundfile as sf
-from omnivoice import OmniVoice
-import tempfile
 import os
+import subprocess
+import zipfile
+import requests
+from huggingface_hub import hf_hub_download, login
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+hf_token = os.getenv("HF_TOKEN")
+login(token=hf_token)
 
-model = OmniVoice.from_pretrained(
-    "k2-fsa/OmniVoice",
-    device_map="cuda:0" if device == "cuda" else "cpu",
-    dtype=torch.float16 if device == "cuda" else torch.float32
-)
+BASE_DIR = os.getcwd()
+SDCPP_DIR = os.path.join(BASE_DIR, "sdcpp")
+
+# =========================
+# 1. دانلود stable-diffusion.cpp
+# =========================
+def setup_sdcpp():
+    if not os.path.exists(SDCPP_DIR):
+        os.makedirs(SDCPP_DIR, exist_ok=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")
+
+        print("Downloading stable-diffusion.cpp...")
+        r = requests.get(zip_url)
+        with open(zip_path, "wb") as f:
+            f.write(r.content)
+
+        print("Extracting...")
+        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+            zip_ref.extractall(SDCPP_DIR)
+
+        os.remove(zip_path)
+
+        # chmod
+        subprocess.run(["chmod", "+x", f"{SDCPP_DIR}/sd-cli"])
+
+
+# =========================
+# 2. دانلود مدل‌ها
+# =========================
+def setup_models():
+    print("Downloading models...")
+
+    model_path = hf_hub_download(
+        repo_id="unsloth/Z-Image-Turbo-GGUF",
+        filename="z-image-turbo-Q4_K_M.gguf"
+    )
 
-def build_voice_prompt(gender, age, pitch, style):
-    attrs = []
-
-    if gender:
-        attrs.append(gender)
-    if age:
-        attrs.append(age)
-    if pitch:
-        attrs.append(pitch)
-    if style:
-        attrs.append(style)
-
-    if len(attrs) > 0:
-        print(attrs)
-        voice_desc = ", ".join(attrs)
-        return voice_desc
-    else:
-        return None
-
-
-def generate(text, ref_audio, gender, age, pitch, style, num_steps, speed):
-
-    output_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
-
-    instruct = build_voice_prompt(gender, age, pitch, style)
-
-    if ref_audio is not None:
-        audio = model.generate(
-            text=text,
-            language='Persian',
-            instruct=instruct,
-            ref_audio=ref_audio,
-            num_step=int(num_steps),
-            speed=float(speed)
-        )
-    else:
-        audio = model.generate(
-            text=text,
-            language='Persian',
-            instruct=instruct,
-            num_step=int(num_steps),
-            speed=float(speed)
-        )
-
-    sf.write(output_path, audio[0], 24000)
+    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-2507-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="Text"),
-
-        gr.Audio(type="filepath", label="Reference Voice (optional)"),
-
-        gr.Dropdown(
-            choices=["male", "female"],
-            label="Gender",
-            value=None
-        ),
-
-        gr.Dropdown(
-            choices=["child", "teenager", "young adult", "middle-aged", "elderly"],
-            label="Age",
-            value=None
-        ),
-
-        gr.Dropdown(
-            choices=["very low pitch", "low pitch", "medium pitch", "high pitch"],
-            label="Pitch",
-            value=None
-        ),
-
-        gr.Dropdown(
-            choices=["normal", "whisper", "calm", "angry"],
-            label="Style",
-            value=None
-        ),
-
-        gr.Slider(
-            minimum=1,
-            maximum=32,
-            value=4,
-            step=1,
-            label="num_steps"
-        ),
-
-        gr.Slider(
-            minimum=0.5,
-            maximum=1.5,
-            value=0.8,
-            step=0.05,
-            label="Speech Speed"
-        )
-    ],
-
-    outputs=gr.Audio(type="filepath"),
-
-    title="OmniVoice TTS (Voice Design + Cloning)",
-    description="Control voice with gender, age, pitch, style + num_steps"
+    inputs=gr.Textbox(label="Prompt"),
+    outputs=gr.Image(label="Generated Image"),
+    title="Z-Image Turbo GGUF (CPU)"
 )
 
 demo.launch()