app.py

import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, LCMScheduler
import time
import json
import hashlib
import base64
import os
import io
import traceback
import threading
from datetime import datetime, timedelta
from PIL import Image

MODEL_ID = "SimianLuo/LCM_Dreamshaper_v7"
DEVICE = "cpu"
CACHE_DIR = "/tmp/generated_images"
CACHE_TTL_SECONDS = 86400
GENERATION_TIMEOUT = 120

os.makedirs(CACHE_DIR, exist_ok=True)

print("Loading Advanced Quality Stack (V2) - MicroCore Studio Edition...")

pipe = StableDiffusionPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,
    safety_checker=None,
    requires_safety_checker=False
)
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.2, b2=1.4)
pipe_i2i = StableDiffusionImg2ImgPipeline(**pipe.components)
pipe.enable_attention_slicing(1)
pipe_i2i.enable_attention_slicing(1)

print("Models loaded and FreeU enabled.")

STYLES = {
    "None": "{prompt}",
    "Cinematic": "cinematic photo, {prompt}, highly detailed, 8k, sharp focus, dramatic lighting, film grain",
    "Photorealistic": "professional photography, {prompt}, realistic, 35mm lens, f/1.8, bokeh, masterpiece, ultra-detailed",
    "Digital Art": "digital painting, {prompt}, intricate detail, vibrant colors, fantasy art, smooth, sharp edges",
    "Cyberpunk": "cyberpunk style, {prompt}, neon lights, rainy street, high contrast, futuristic, blade runner aesthetic",
    "Anime": "anime style, {prompt}, hand-drawn, high resolution, vibrant, clean lines"
}

class CircuitBreaker:
    def __init__(self, failure_threshold=3, reset_timeout=60):
        self.failure_threshold = failure_threshold
        self.reset_timeout = reset_timeout
        self.failure_count = 0
        self.last_failure_time = None
        self.state = "CLOSED"
        self.lock = threading.Lock()

    def is_open(self):
        with self.lock:
            if self.state == "OPEN":
                if time.time() - self.last_failure_time >= self.reset_timeout:
                    self.state = "HALF_OPEN"
                    print("[CircuitBreaker] Transitioning to HALF_OPEN")
                    return False
                return True
            return False

    def record_success(self):
        with self.lock:
            self.failure_count = 0
            self.state = "CLOSED"

    def record_failure(self):
        with self.lock:
            self.failure_count += 1
            self.last_failure_time = time.time()
            if self.failure_count >= self.failure_threshold:
                self.state = "OPEN"
                print(f"[CircuitBreaker] OPEN after {self.failure_count} failures, pausing for {self.reset_timeout}s")

circuit_breaker = CircuitBreaker(failure_threshold=3, reset_timeout=60)

def compute_cache_key(payload_dict):
    payload_str = json.dumps(payload_dict, sort_keys=True)
    return hashlib.sha256(payload_str.encode()).hexdigest()

def get_cached_image(cache_key):
    cache_path = os.path.join(CACHE_DIR, f"{cache_key}.png")
    if not os.path.exists(cache_path):
        return None
    file_mtime = datetime.fromtimestamp(os.path.getmtime(cache_path))
    if datetime.now() - file_mtime > timedelta(seconds=CACHE_TTL_SECONDS):
        try:
            os.remove(cache_path)
        except OSError:
            pass
        return None
    try:
        with open(cache_path, "rb") as f:
            return f.read()
    except IOError:
        return None

def save_cached_image(cache_key, image_data):
    cache_path = os.path.join(CACHE_DIR, f"{cache_key}.png")
    try:
        with open(cache_path, "wb") as f:
            f.write(image_data)
    except IOError as e:
        print(f"[Cache] Warning: Could not save to {cache_path}: {e}")

def image_to_base64(image):
    buf = io.BytesIO()
    image.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode("utf-8")

def generate_advanced(prompt, negative_prompt, style, steps, guidance, polish_intensity):
    if circuit_breaker.is_open():
        error_msg = json.dumps({"error": "Service temporarily unavailable due to high error rate. Please retry in 60s."})
        raise gr.Error(error_msg)

    if not prompt or not isinstance(prompt, str) or len(prompt.strip()) == 0:
        error_msg = json.dumps({"error": "Invalid payload: 'prompt' is required and must be a non-empty string.", "code": 400})
        raise gr.Error(error_msg)

    try:
        steps = int(steps) if steps is not None else 8
        guidance = float(guidance) if guidance is not None else 1.5
        polish_intensity = float(polish_intensity) if polish_intensity is not None else 0.3
    except (ValueError, TypeError):
        error_msg = json.dumps({"error": "Invalid payload: steps, guidance, and polish_intensity must be numeric.", "code": 400})
        raise gr.Error(error_msg)

    cache_payload = {
        "prompt": prompt,
        "negative_prompt": negative_prompt or "",
        "style": style,
        "steps": steps,
        "guidance": guidance,
        "polish_intensity": polish_intensity
    }
    cache_key = compute_cache_key(cache_payload)

    cached = get_cached_image(cache_key)
    if cached is not None:
        print(f"[Cache] HIT for key {cache_key[:16]}...")
        cached_image = Image.open(io.BytesIO(cached))
        b64_data = base64.b64encode(cached).decode("utf-8")
        return cached_image, f"CACHED (TTL: 24h) | Cache Key: {cache_key[:16]}..."

    start_time = time.time()

    try:
        full_prompt = STYLES[style].format(prompt=prompt)
    except KeyError:
        error_msg = json.dumps({"error": f"Invalid payload: unknown style '{style}'. Valid styles: {list(STYLES.keys())}", "code": 400})
        raise gr.Error(error_msg)

    print(f"Stage 1: Generating base image with {style} style...")
    try:
        base_image = pipe(
            prompt=full_prompt,
            negative_prompt=negative_prompt if negative_prompt else "blurry, low quality, distorted",
            num_inference_steps=steps,
            guidance_scale=guidance,
            width=768,
            height=768,
        ).images[0]
    except Exception as e:
        circuit_breaker.record_failure()
        print(f"[Pipeline] Stage 1 failed: {e}")
        error_msg = json.dumps({"error": f"Generation failed in stage 1: {str(e)}", "code": 500})
        raise gr.Error(error_msg)

    if polish_intensity > 0:
        print(f"Stage 2: Applying Polish Pass (Intensity: {polish_intensity})...")
        try:
            refined_image = pipe_i2i(
                prompt=full_prompt,
                negative_prompt=negative_prompt,
                image=base_image,
                strength=polish_intensity,
                num_inference_steps=max(1, int(steps / 2)),
                guidance_scale=guidance,
            ).images[0]
            final_image = refined_image
        except Exception as e:
            print(f"[Pipeline] Stage 2 (polish) failed, returning base image: {e}")
            final_image = base_image
    else:
        final_image = base_image

    duration = round(time.time() - start_time, 2)
    circuit_breaker.record_success()

    try:
        img_bytes = image_to_base64(final_image)
        raw_bytes = base64.b64decode(img_bytes)
        save_cached_image(cache_key, raw_bytes)
        print(f"[Cache] SAVED key {cache_key[:16]}... ({len(raw_bytes)} bytes)")
    except Exception as e:
        print(f"[Cache] Warning: Failed to save cache: {e}")

    return final_image, f"Quality Optimized Generation in {duration}s | Cache Key: {cache_key[:16]}..."

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚀 Advanced CPU Image Gen V2 — MicroCore Studio Edition")
    gr.Markdown("Using **FreeU** + **Two-Stage Refinement** + **Caching** + **Circuit Breaker** for reliable API integration.")

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", placeholder="Describe your vision...", lines=3)
            style = gr.Dropdown(choices=list(STYLES.keys()), value="Photorealistic", label="Style Engine (Auto-Boosting)")

            with gr.Accordion("Advanced Settings", open=False):
                negative_prompt = gr.Textbox(label="Negative Prompt", value="blurry, ugly, low quality, deformed")
                steps = gr.Slider(4, 12, value=8, step=1, label="Steps")
                guidance = gr.Slider(1.0, 4.0, value=1.5, step=0.1, label="Guidance Scale")
                polish = gr.Slider(0.0, 0.5, value=0.3, step=0.05, label="Polish Intensity (Refiner Pass)")

            btn = gr.Button("🎨 Generate High Quality Image", variant="primary")

        with gr.Column():
            output_image = gr.Image(label="V2 Advanced Output", format="png")
            status = gr.Text(label="Engine Status")

    btn.click(
        fn=generate_advanced,
        inputs=[prompt, negative_prompt, style, steps, guidance, polish],
        outputs=[output_image, status]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)