Spaces:

Hfzune
/

Test

Sleeping

Test / app.py

Update app.py

df67f98 verified 3 days ago

9.27 kB

import os

# Must be set before torch/diffusers import
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

# MUST come before torch / diffusers (ZeroGPU scanner)
import spaces

import torch
import gradio as gr
import asyncio
import httpx
import tempfile
import traceback

from fastapi import Request
from fastapi.responses import JSONResponse
from starlette.routing import Route as StarletteRoute

# ✅ CORRECT pipeline for Juggernaut Z — NOT Lumina2Pipeline
from diffusers import ZImagePipeline


# ============================================================
# CONFIG
# ============================================================
MODEL_ID = os.getenv(
    "MODEL_ID",
    "RunDiffusion/Juggernaut-Z-Image",
)
TELEGRAM_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "")
TELEGRAM_API = f"https://api.telegram.org/bot{TELEGRAM_TOKEN}" if TELEGRAM_TOKEN else ""
MAX_IMAGES = int(os.getenv("MAX_IMAGES_PER_REQUEST", "10"))
GPU_DURATION = int(os.getenv("GPU_DURATION_SEC", "120"))

# Juggernaut Z recommended settings (official model card)
STEPS = int(os.getenv("GEN_STEPS", "35"))
GUIDANCE = float(os.getenv("GEN_GUIDANCE", "6.0"))
GEN_SIZE = int(os.getenv("GEN_SIZE", "1024"))


# ============================================================
# MODEL LOAD (module level — ZeroGPU streams weights on first GPU call)
# ============================================================
print(f"[BOOT] Loading Juggernaut Z pipeline: {MODEL_ID}")

# Use ZImagePipeline. If your diffusers version rejects it, fall back to:
# from diffusers import DiffusionPipeline
# pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
pipe = ZImagePipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
)

print("[BOOT] Pipeline loaded (weights on CPU/disk; ZeroGPU will stream to VRAM)")


# ============================================================
# TELEGRAM HELPERS
# ============================================================
async def send_message(chat_id: int, text: str) -> None:
    if not TELEGRAM_TOKEN:
        return
    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            await client.post(
                f"{TELEGRAM_API}/sendMessage",
                json={
                    "chat_id": chat_id,
                    "text": text,
                    "parse_mode": "HTML",
                    "disable_web_page_preview": True,
                },
            )
    except Exception as e:
        print(f"[TG] send_message failed: {e}")


async def send_photo(chat_id: int, path: str, caption: str = "") -> None:
    if not TELEGRAM_TOKEN:
        return
    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            with open(path, "rb") as f:
                await client.post(
                    f"{TELEGRAM_API}/sendPhoto",
                    data={"chat_id": chat_id, "caption": caption, "parse_mode": "HTML"},
                    files={"photo": f},
                )
    except Exception as e:
        print(f"[TG] send_photo failed: {e}")


# ============================================================
# GPU FUNCTION — the ONLY @spaces.GPU function in this Space
# ============================================================
@spaces.GPU(duration=GPU_DURATION)
def generate_image_gpu(prompt: str, seed: int = -1) -> str:
    """
    Generate one hyper-realistic image with Juggernaut Z.
    Official recommended settings: 35 steps, CFG 6.0.
    """
    # Move to GPU on first call. ZeroGPU caches weights in VRAM afterwards.
    if pipe.device.type != "cuda":
        pipe.to("cuda")

    generator = None
    if seed != -1:
        generator = torch.Generator(device="cuda").manual_seed(int(seed))

    result = pipe(
        prompt=prompt,
        generator=generator,
        num_inference_steps=STEPS,
        guidance_scale=GUIDANCE,
        height=GEN_SIZE,
        width=GEN_SIZE,
    )

    image = result.images[0]

    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
    image.save(tmp)
    return tmp


# ============================================================
# TASK QUEUE + WORKER
# ============================================================
task_queue: asyncio.Queue = asyncio.Queue()
_worker_task: asyncio.Task | None = None


async def worker() -> None:
    print("[WORKER] started")
    while True:
        chat_id = None
        try:
            chat_id, prompt, num = await task_queue.get()
            await send_message(chat_id, f"🎨 Generating {num} hyper-realistic image(s)...")
            for i in range(num):
                await send_message(chat_id, f"⏳ Image {i + 1}/{num} — rendering (35 steps)...")
                try:
                    # Run the sync GPU function in a thread so we don't block the loop
                    path = await asyncio.to_thread(generate_image_gpu, prompt)
                    await send_photo(chat_id, path, f"✅ Image {i + 1}/{num}")
                    try:
                        os.unlink(path)
                    except OSError:
                        pass
                except Exception as e:
                    print(f"[WORKER] image {i + 1} failed: {e}\n{traceback.format_exc()}")
                    await send_message(chat_id, f"❌ Image {i + 1} failed: {e}")
            await send_message(chat_id, "🎉 All done!")
        except Exception as e:
            print(f"[WORKER] fatal: {e}\n{traceback.format_exc()}")
            if chat_id:
                try:
                    await send_message(chat_id, f"❌ Worker error: {e}")
                except Exception:
                    pass
        finally:
            task_queue.task_done()


# ============================================================
# FASTAPI ROUTES
# ============================================================
async def telegram_webhook(request: Request):
    global _worker_task
    # Lazy-start the worker inside the running event loop
    if _worker_task is None:
        _worker_task = asyncio.create_task(worker())

    try:
        update = await request.json()
        msg = update.get("message", {})
        chat_id = msg.get("chat", {}).get("id")
        text = msg.get("text", "")

        if chat_id and text:
            if text.startswith("/start"):
                await send_message(
                    chat_id,
                    "👋 <b>Hyper-Realistic Image Bot (Juggernaut Z)</b>\n\n"
                    f"Send <code>/imagine N &lt;prompt&gt;</code> (N = 1..{MAX_IMAGES}).\n"
                    "Example: <code>/imagine 2 cinematic portrait of an old fisherman, dramatic lighting</code>",
                )
            elif text.startswith("/imagine"):
                parts = text.split(maxsplit=2)
                if len(parts) >= 3:
                    try:
                        num = int(parts[1])
                        if not 1 <= num <= MAX_IMAGES:
                            raise ValueError
                    except ValueError:
                        await send_message(chat_id, f"❌ Number must be 1..{MAX_IMAGES}")
                        return JSONResponse({"ok": True})
                    prompt = parts[2]
                    await task_queue.put((chat_id, prompt, num))
                    await send_message(chat_id, f"✅ Queued {num} image(s): <i>{prompt}</i>")
                else:
                    await send_message(chat_id, "Usage: <code>/imagine N prompt</code>")
    except Exception as e:
        print(f"[WEBHOOK] error: {e}\n{traceback.format_exc()}")

    return JSONResponse({"ok": True})


async def health_check(request: Request):
    return JSONResponse({
        "ok": True,
        "status": "running",
        "queue_size": task_queue.qsize(),
        "worker_alive": _worker_task is not None and not _worker_task.done(),
    })


# ============================================================
# GRADIO DEMO — hidden button satisfies the ZeroGPU scanner
# ============================================================
with gr.Blocks() as demo:
    gr.Markdown("# 🎨 Hyper-Realistic Image Bot (Juggernaut Z)\n\nServes a Telegram bot via `/telegram`.")
    _btn = gr.Button("init-gpu", visible=False)
    _in = gr.Textbox(value="cinematic portrait, dramatic lighting, photorealistic", visible=False)
    _out = gr.File(visible=False)
    _btn.click(fn=generate_image_gpu, inputs=[_in], outputs=[_out])


# ============================================================
# LAUNCH — proven pattern from your history
# ============================================================
demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    share=False,
    ssr_mode=False,             # prevents Node proxy from swallowing /telegram
    prevent_thread_lock=True,   # returns control so we can inject routes
)

# Inject routes at index 0 so they beat Gradio's catch-all Mount
live_app = demo.app
live_app.router.routes.insert(
    0, StarletteRoute("/telegram", telegram_webhook, methods=["POST"])
)
live_app.router.routes.insert(
    0, StarletteRoute("/health", health_check, methods=["GET"])
)
live_app.build_middleware_stack()

print("[BOOT] Routes injected: POST /telegram, GET /health")
print("[BOOT] Set webhook → https://<space>.hf.space/telegram")

demo.block_thread()

	import os

	# Must be set before torch/diffusers import
	os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

	# MUST come before torch / diffusers (ZeroGPU scanner)
	import spaces

	import torch
	import gradio as gr
	import asyncio
	import httpx
	import tempfile
	import traceback

	from fastapi import Request
	from fastapi.responses import JSONResponse
	from starlette.routing import Route as StarletteRoute

	# ✅ CORRECT pipeline for Juggernaut Z — NOT Lumina2Pipeline
	from diffusers import ZImagePipeline


	# ============================================================
	# CONFIG
	# ============================================================
	MODEL_ID = os.getenv(
	"MODEL_ID",
	"RunDiffusion/Juggernaut-Z-Image",
	)
	TELEGRAM_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "")
	TELEGRAM_API = f"https://api.telegram.org/bot{TELEGRAM_TOKEN}" if TELEGRAM_TOKEN else ""
	MAX_IMAGES = int(os.getenv("MAX_IMAGES_PER_REQUEST", "10"))
	GPU_DURATION = int(os.getenv("GPU_DURATION_SEC", "120"))

	# Juggernaut Z recommended settings (official model card)
	STEPS = int(os.getenv("GEN_STEPS", "35"))
	GUIDANCE = float(os.getenv("GEN_GUIDANCE", "6.0"))
	GEN_SIZE = int(os.getenv("GEN_SIZE", "1024"))


	# ============================================================
	# MODEL LOAD (module level — ZeroGPU streams weights on first GPU call)
	# ============================================================
	print(f"[BOOT] Loading Juggernaut Z pipeline: {MODEL_ID}")

	# Use ZImagePipeline. If your diffusers version rejects it, fall back to:
	# from diffusers import DiffusionPipeline
	# pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
	pipe = ZImagePipeline.from_pretrained(
	MODEL_ID,
	torch_dtype=torch.bfloat16,
	)

	print("[BOOT] Pipeline loaded (weights on CPU/disk; ZeroGPU will stream to VRAM)")


	# ============================================================
	# TELEGRAM HELPERS
	# ============================================================
	async def send_message(chat_id: int, text: str) -> None:
	if not TELEGRAM_TOKEN:
	return
	try:
	async with httpx.AsyncClient(timeout=15.0) as client:
	await client.post(
	f"{TELEGRAM_API}/sendMessage",
	json={
	"chat_id": chat_id,
	"text": text,
	"parse_mode": "HTML",
	"disable_web_page_preview": True,
	},
	)
	except Exception as e:
	print(f"[TG] send_message failed: {e}")


	async def send_photo(chat_id: int, path: str, caption: str = "") -> None:
	if not TELEGRAM_TOKEN:
	return
	try:
	async with httpx.AsyncClient(timeout=120.0) as client:
	with open(path, "rb") as f:
	await client.post(
	f"{TELEGRAM_API}/sendPhoto",
	data={"chat_id": chat_id, "caption": caption, "parse_mode": "HTML"},
	files={"photo": f},
	)
	except Exception as e:
	print(f"[TG] send_photo failed: {e}")


	# ============================================================
	# GPU FUNCTION — the ONLY @spaces.GPU function in this Space
	# ============================================================
	@spaces.GPU(duration=GPU_DURATION)
	def generate_image_gpu(prompt: str, seed: int = -1) -> str:
	"""
	Generate one hyper-realistic image with Juggernaut Z.
	Official recommended settings: 35 steps, CFG 6.0.
	"""
	# Move to GPU on first call. ZeroGPU caches weights in VRAM afterwards.
	if pipe.device.type != "cuda":
	pipe.to("cuda")

	generator = None
	if seed != -1:
	generator = torch.Generator(device="cuda").manual_seed(int(seed))

	result = pipe(
	prompt=prompt,
	generator=generator,
	num_inference_steps=STEPS,
	guidance_scale=GUIDANCE,
	height=GEN_SIZE,
	width=GEN_SIZE,
	)

	image = result.images[0]

	tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
	image.save(tmp)
	return tmp


	# ============================================================
	# TASK QUEUE + WORKER
	# ============================================================
	task_queue: asyncio.Queue = asyncio.Queue()
	_worker_task: asyncio.Task \| None = None


	async def worker() -> None:
	print("[WORKER] started")
	while True:
	chat_id = None
	try:
	chat_id, prompt, num = await task_queue.get()
	await send_message(chat_id, f"🎨 Generating {num} hyper-realistic image(s)...")
	for i in range(num):
	await send_message(chat_id, f"⏳ Image {i + 1}/{num} — rendering (35 steps)...")
	try:
	# Run the sync GPU function in a thread so we don't block the loop
	path = await asyncio.to_thread(generate_image_gpu, prompt)
	await send_photo(chat_id, path, f"✅ Image {i + 1}/{num}")
	try:
	os.unlink(path)
	except OSError:
	pass
	except Exception as e:
	print(f"[WORKER] image {i + 1} failed: {e}\n{traceback.format_exc()}")
	await send_message(chat_id, f"❌ Image {i + 1} failed: {e}")
	await send_message(chat_id, "🎉 All done!")
	except Exception as e:
	print(f"[WORKER] fatal: {e}\n{traceback.format_exc()}")
	if chat_id:
	try:
	await send_message(chat_id, f"❌ Worker error: {e}")
	except Exception:
	pass
	finally:
	task_queue.task_done()


	# ============================================================
	# FASTAPI ROUTES
	# ============================================================
	async def telegram_webhook(request: Request):
	global _worker_task
	# Lazy-start the worker inside the running event loop
	if _worker_task is None:
	_worker_task = asyncio.create_task(worker())

	try:
	update = await request.json()
	msg = update.get("message", {})
	chat_id = msg.get("chat", {}).get("id")
	text = msg.get("text", "")

	if chat_id and text:
	if text.startswith("/start"):
	await send_message(
	chat_id,
	"👋 <b>Hyper-Realistic Image Bot (Juggernaut Z)</b>\n\n"
	f"Send <code>/imagine N <prompt></code> (N = 1..{MAX_IMAGES}).\n"
	"Example: <code>/imagine 2 cinematic portrait of an old fisherman, dramatic lighting</code>",
	)
	elif text.startswith("/imagine"):
	parts = text.split(maxsplit=2)
	if len(parts) >= 3:
	try:
	num = int(parts[1])
	if not 1 <= num <= MAX_IMAGES:
	raise ValueError
	except ValueError:
	await send_message(chat_id, f"❌ Number must be 1..{MAX_IMAGES}")
	return JSONResponse({"ok": True})
	prompt = parts[2]
	await task_queue.put((chat_id, prompt, num))
	await send_message(chat_id, f"✅ Queued {num} image(s): <i>{prompt}</i>")
	else:
	await send_message(chat_id, "Usage: <code>/imagine N prompt</code>")
	except Exception as e:
	print(f"[WEBHOOK] error: {e}\n{traceback.format_exc()}")

	return JSONResponse({"ok": True})


	async def health_check(request: Request):
	return JSONResponse({
	"ok": True,
	"status": "running",
	"queue_size": task_queue.qsize(),
	"worker_alive": _worker_task is not None and not _worker_task.done(),
	})


	# ============================================================
	# GRADIO DEMO — hidden button satisfies the ZeroGPU scanner
	# ============================================================
	with gr.Blocks() as demo:
	gr.Markdown("# 🎨 Hyper-Realistic Image Bot (Juggernaut Z)\n\nServes a Telegram bot via `/telegram`.")
	_btn = gr.Button("init-gpu", visible=False)
	_in = gr.Textbox(value="cinematic portrait, dramatic lighting, photorealistic", visible=False)
	_out = gr.File(visible=False)
	_btn.click(fn=generate_image_gpu, inputs=[_in], outputs=[_out])


	# ============================================================
	# LAUNCH — proven pattern from your history
	# ============================================================
	demo.launch(
	server_name="0.0.0.0",
	server_port=7860,
	share=False,
	ssr_mode=False, # prevents Node proxy from swallowing /telegram
	prevent_thread_lock=True, # returns control so we can inject routes
	)

	# Inject routes at index 0 so they beat Gradio's catch-all Mount
	live_app = demo.app
	live_app.router.routes.insert(
	0, StarletteRoute("/telegram", telegram_webhook, methods=["POST"])
	)
	live_app.router.routes.insert(
	0, StarletteRoute("/health", health_check, methods=["GET"])
	)
	live_app.build_middleware_stack()

	print("[BOOT] Routes injected: POST /telegram, GET /health")
	print("[BOOT] Set webhook → https://<space>.hf.space/telegram")

	demo.block_thread()