app.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
server_open_image_ai.py
Gradio-based, server-ready, open-source image AI:
- No paid APIs; uses local diffusers (SDXL/SD1.5) if available.
- Offline optical-illusion fallback.
- Persian error messages with actionable fixes.
- Prompt variation, post-process filters, caching, simple logging.

Run:
  python server_open_image_ai.py
Then open the printed local URL or host=0.0.0.0 for external access.
"""

import os
import sys
import time
import json
import hashlib
import random
from datetime import datetime
from pathlib import Path

# Lazy install open-source libs only
def _lazy_install(pkg):
    try:
        __import__(pkg.split("==")[0])
    except ImportError:
        os.system(f"{sys.executable} -m pip install --quiet {pkg}")

for pkg in ("gradio>=4.38.0", "torch", "diffusers", "accelerate", "safetensors", "Pillow"):
    _lazy_install(pkg)

import gradio as gr
import torch
from PIL import Image, ImageOps, ImageFilter, ImageDraw
from diffusers import StableDiffusionPipeline, DiffusionPipeline

# ---------------------------
# Paths & config
# ---------------------------
APP_NAME = "Open Image AI (Server)"
ROOT_DIR = Path(__file__).parent.resolve()
OUTPUT_DIR = Path(os.environ.get("OIA_OUT", ROOT_DIR / "outputs")).resolve()
CACHE_DIR = OUTPUT_DIR / "cache"
META_DIR = OUTPUT_DIR / "meta"
for d in (OUTPUT_DIR, CACHE_DIR, META_DIR):
    d.mkdir(parents=True, exist_ok=True)

DEFAULT_STEPS = int(os.environ.get("OIA_STEPS", "30"))
DEFAULT_GUIDANCE = float(os.environ.get("OIA_GUIDANCE", "7.5"))
DEFAULT_SIZE = int(os.environ.get("OIA_SIZE", "768"))

OPEN_MODELS = {
    "SDXL (base)": "stabilityai/stable-diffusion-xl-base-1.0",
    "SD 1.5 (classic)": "runwayml/stable-diffusion-v1-5",
}

# ---------------------------
# Prompt helpers
# ---------------------------
OPTICAL_ILLUSION_PRESETS = [
    "high-contrast black and white optical illusion, impossible geometry, Penrose stairs, concentric lines, parallax shifts",
    "moiré patterns, nested grids, Escher-inspired recursion, tilt-shift depth cues, perspective ambiguity",
    "rotational symmetry, non-Euclidean corridor, variable line thickness, alternating diagonals, dizzying parallax",
    "spiral tunnels, zigzag pathways, impossible cube, nested frames, sharp edges, stark contrast",
]
WEIGHTS = [
    "hyper-detailed", "vector-sharp edges", "ultra high resolution",
    "minimal noise", "geometric precision", "dynamic parallax",
    "photorealistic lighting"
]
MODS = [
    "isometric view", "45-degree inclination", "multi-angle perception",
    "viewer-dependent illusion", "top-down perspective"
]

def vary_prompt(base: str) -> str:
    extra = ", ".join(random.sample(OPTICAL_ILLUSION_PRESETS, k=random.randint(1, 3)))
    hint = ", ".join(random.sample(WEIGHTS, k=random.randint(2, 4)))
    mod = random.choice(MODS)
    return f"{base}, {extra}, {hint}, {mod}"

def prompt_hash(p: str) -> str:
    return hashlib.sha256(p.encode("utf-8")).hexdigest()[:16]

def cache_path_for(prompt: str) -> Path:
    return CACHE_DIR / f"{prompt_hash(prompt)}.png"

def write_meta(meta: dict, fname: Path):
    meta["timestamp"] = datetime.utcnow().isoformat()
    with open(META_DIR / (fname.stem + ".json"), "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False, indent=2)

# ---------------------------
# Offline optical-illusion generator
# ---------------------------
def local_illusion(size=DEFAULT_SIZE) -> Image.Image:
    img = Image.new("RGB", (size, size), "white")
    draw = ImageDraw.Draw(img)
    step = 6
    for r in range(24, size//2, step):
        val = 0 if (r//step) % 2 == 0 else 255
        for x in range(size):
            if abs(x - size//2) == r:
                draw.line([(x, 0), (x, size)], fill=(val, val, val))
        for y in range(size):
            if abs(y - size//2) == r:
                draw.line([(0, y), (size, y)], fill=(val, val, val))
    hatch = Image.new("L", (size, size), color=255)
    hpx = hatch.load()
    for x in range(size):
        for y in range(size):
            if (x + y) % 18 == 0:
                hpx[x, y] = 0
    merged = Image.merge("RGB", (ImageOps.autocontrast(hatch),) * 3)
    merged = Image.blend(img, merged, 0.35).filter(ImageFilter.SHARPEN)
    merged = merged.rotate(random.choice([4, -6, 8, -10]), resample=Image.BICUBIC, expand=False)
    return merged

# ---------------------------
# Post-process filters
# ---------------------------
POST_MODES = [
    "none", "illusion_boost", "bw_high", "moire_mix", "invert", "edge_halo",
    "soft_glow", "posterize4", "contrast_max", "rotate_slight", "grain_fine"
]

def post_process(img: Image.Image, mode: str) -> Image.Image:
    if mode == "illusion_boost":
        return ImageOps.autocontrast(img).filter(ImageFilter.UnsharpMask(radius=2, percent=180, threshold=3))
    if mode == "bw_high":
        return ImageOps.autocontrast(ImageOps.grayscale(img)).convert("RGB")
    if mode == "moire_mix":
        return ImageOps.posterize(img.filter(ImageFilter.DETAIL).filter(ImageFilter.SHARPEN), bits=4)
    if mode == "invert":
        return ImageOps.invert(img)
    if mode == "edge_halo":
        return img.filter(ImageFilter.UnsharpMask(radius=3, percent=240, threshold=2))
    if mode == "soft_glow":
        blur = img.filter(ImageFilter.GaussianBlur(radius=2))
        return Image.blend(img, blur, alpha=0.3)
    if mode == "posterize4":
        return ImageOps.posterize(img, bits=4)
    if mode == "contrast_max":
        return ImageOps.autocontrast(img)
    if mode == "rotate_slight":
        return img.rotate(random.choice([-3, 3, -2, 2]), resample=Image.BICUBIC, expand=False)
    if mode == "grain_fine":
        return img.filter(ImageFilter.SMOOTH_MORE)
    return img

# ---------------------------
# Model hub (diffusers)
# ---------------------------
class ModelHub:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.pipe = None
        self.model_id = None

    def load(self, model_id: str, fp16=True) -> tuple[bool, str]:
        self.model_id = model_id
        dtype = torch.float16 if (fp16 and self.device == "cuda") else torch.float32
        try:
            if "stable-diffusion-xl" in model_id:
                self.pipe = DiffusionPipeline.from_pretrained(
                    model_id, torch_dtype=dtype,
                    variant="fp16" if dtype == torch.float16 else None
                )
            else:
                self.pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype)
            self.pipe = self.pipe.to(self.device)
            if self.device == "cuda":
                # حافظه بهتر برای GPUهای متوسط
                self.pipe.enable_attention_slicing()
            return True, "مدل با موفقیت بارگذاری شد."
        except Exception as e:
            self.pipe = None
            return False, f"خطا در بارگذاری مدل: {e}\nراه‌حل: مطمئن شوید فضای دیسک کافی است و اینترنت یا فایل‌های مدل در دسترس‌اند. اگر GPU ندارید، اجرای CPU ممکن است کند باشد."

    def generate(self, prompt: str, steps: int, guidance: float, seed: int | None, height=DEFAULT_SIZE, width=DEFAULT_SIZE) -> tuple[Image.Image, str]:
        if self.pipe is None:
            return local_illusion(size=min(height, width)), "مدل در دسترس نیست؛ از مولد آفلاین خطای دید استفاده شد."
        # تعیین seed
        generator = torch.Generator(device=self.device)
        if seed is not None and seed >= 0:
            generator = generator.manual_seed(seed)
        try:
            out = self.pipe(
                prompt,
                height=height, width=width,
                num_inference_steps=steps,
                guidance_scale=guidance,
                generator=generator
            )
            return out.images[0], "تولید موفق."
        except Exception as e:
            return local_illusion(size=min(height, width)), f"خطا در تولید: {e}\nراه‌حل: اندازه تصویر را کاهش دهید، تعداد گام‌ها را کم کنید، یا مدل سبک‌تر (SD 1.5) را انتخاب کنید."

hub = ModelHub()

# ---------------------------
# Core generate function (for Gradio)
# ---------------------------
def generate_ui(prompt, model_name, steps, guidance, size, seed, vary, post_mode):
    # آماده‌سازی پرومپت
    base_prompt = prompt.strip()
    if not base_prompt:
        return None, "پرومپت خالی است.\nراه‌حل: یک توضیح کوتاه یا بلند برای تصویر وارد کنید.", ""

    use_prompt = vary_prompt(base_prompt) if vary else base_prompt

    # کش
    cp = cache_path_for(use_prompt)
    if cp.exists():
        try:
            img = Image.open(cp).convert("RGB")
            img = post_process(img, post_mode)
            fname = OUTPUT_DIR / f"{datetime.utcnow().strftime('%Y%m%d_%H%M%S_%f')}_{prompt_hash(use_prompt)}.png"
            img.save(fname, "PNG", optimize=True)
            write_meta({"prompt": use_prompt, "model": hub.model_id or "local_illusion", "post_mode": post_mode}, fname)
            return fname, "بارگیری از کش.", use_prompt
        except Exception as e:
            # اگر کش خراب بود، ادامه تولید
            pass

    # بارگذاری مدل در صورت نیاز
    model_id = OPEN_MODELS.get(model_name, "")
    if hub.pipe is None or hub.model_id != model_id:
        ok, msg = hub.load(model_id)
        if not ok:
            img = post_process(local_illusion(size=size), post_mode)
            fname = OUTPUT_DIR / f"{datetime.utcnow().strftime('%Y%m%d_%H%M%S_%f')}_fallback_{prompt_hash(use_prompt)}.png"
            img.save(fname, "PNG", optimize=True)
            write_meta({"prompt": use_prompt, "model": "local_illusion", "post_mode": post_mode, "error": msg}, fname)
            return fname, msg, use_prompt

    # تولید
    img, status = hub.generate(use_prompt, steps=steps, guidance=guidance, seed=seed, height=size, width=size)
    img = post_process(img, post_mode)
    fname = OUTPUT_DIR / f"{datetime.utcnow().strftime('%Y%m%d_%H%M%S_%f')}_{prompt_hash(use_prompt)}.png"
    img.save(fname, "PNG", optimize=True)
    write_meta({"prompt": use_prompt, "model": hub.model_id or "local_illusion", "post_mode": post_mode}, fname)
    # ذخیره در کش برای بارهای بعدی
    try:
        cache_path = cache_path_for(use_prompt)
        img.save(cache_path, "PNG", optimize=True)
    except Exception:
        pass
    return fname, status, use_prompt

# ---------------------------
# Gradio UI
# ---------------------------
with gr.Blocks(title=APP_NAME) as demo:
    gr.Markdown(f"# {APP_NAME}\nسامانه وب تصویرساز متن‌باز. هیچ سرویس پولی استفاده نمی‌شود. خطاها با پیام‌های فارسی و راه‌حل برگردانده می‌شوند.")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="پرومپت (توضیح تصویر)", placeholder="مثلاً: پلکان ناممکن با خطوط سیاه و سفید، خطای دید شدید...")
            model_name = gr.Dropdown(choices=list(OPEN_MODELS.keys()), value="SD 1.5 (classic)", label="مدل متن‌باز")
            steps = gr.Slider(5, 100, value=DEFAULT_STEPS, step=1, label="تعداد گام‌ها")
            guidance = gr.Slider(0.0, 20.0, value=DEFAULT_GUIDANCE, step=0.5, label="راهنمایی (Guidance)")
            size = gr.Slider(256, 1024, value=DEFAULT_SIZE, step=64, label="اندازه (پیکسل مربع)")
            seed = gr.Number(value=-1, precision=0, label="Seed (برای تکرارپذیری؛ -1 خودکار)")
            vary = gr.Checkbox(value=True, label="تنوع‌دهی خودکار پرومپت")
            post_mode = gr.Dropdown(choices=POST_MODES, value="none", label="پس‌پردازش")
            btn = gr.Button("تولید تصویر")
        with gr.Column():
            out_file = gr.File(label="خروجی (فایل PNG ذخیره‌شده)")
            out_status = gr.Textbox(label="وضعیت / پیام فارسی", interactive=False)
            out_prompt = gr.Textbox(label="پرومپت نهایی", interactive=False)

    btn.click(
        generate_ui,
        inputs=[prompt, model_name, steps, guidance, size, seed, vary, post_mode],
        outputs=[out_file, out_status, out_prompt]
    )

def main():
    # برای سرورها:
    # share=False برای داخلی؛ اگر می‌خواهید از بیرون دسترسی بدهید، host را 0.0.0.0 تنظیم کنید.
    demo.queue(max_size=32).launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", "7860")), share=False)

if __name__ == "__main__":
    main()