import streamlit as st import io import requests import random import datetime import qrcode import json from PIL import Image, ImageDraw, ImageFont from huggingface_hub import hf_hub_download from deep_translator import GoogleTranslator import numpy as np import os import shutil DATASET_ID = os.getenv("DATASET_ID") HF_TOKEN = os.getenv("HF_TOKEN") try: from streamlit.runtime.scriptrunner import add_script_run_ctx except ImportError: def add_script_run_ctx(t): pass OPENCV_AVAILABLE = False try: import cv2 from cv2 import dnn_superres OPENCV_AVAILABLE = True except ImportError: pass PSUTIL_AVAILABLE = False try: import psutil PSUTIL_AVAILABLE = True except ImportError: pass # [V30.12] 内存优化：使用上下文管理器确保图片句柄立即释放 def generate_thumbnail(image_bytes, size=(360, 360), quality=80): try: with Image.open(io.BytesIO(image_bytes)) as img: if img.mode in ("RGBA", "P"): img = img.convert("RGB") # 使用低算力算法 img.thumbnail(size, Image.Resampling.BILINEAR) buf = io.BytesIO() img.save(buf, format="JPEG", quality=quality, optimize=True) return buf.getvalue() except Exception: return None def get_disk_info(path): try: total_size = 0 if os.path.exists(path): for dirpath, dirnames, filenames in os.walk(path): for f in filenames: fp = os.path.join(dirpath, f) if not os.path.islink(fp): total_size += os.path.getsize(fp) used_gb = total_size / (1024**3) total_host, used_host, free_host = shutil.disk_usage("/") return used_gb, free_host / (1024**3), total_host / (1024**3), (used_host / total_host) * 100 except Exception: return 0, 0, 0, 0 def _download_image_from_hf(repo_id, filename, token): try: # [V30.12] 移除 resume_download path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="dataset", token=token, local_files_only=True) with open(path, "rb") as f: return f.read() except: try: path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="dataset", token=token) with open(path, "rb") as f: return f.read() except Exception: return None def ensure_png_bytes(image_data): try: with Image.open(io.BytesIO(image_data)) as img: buf = io.BytesIO() img.save(buf, format="PNG") return buf.getvalue() except Exception: return image_data def auto_translate(text): if not text: return text, False if any("\u4e00" <= char <= "\u9fff" for char in text): try: return GoogleTranslator(source='auto', target='en').translate(text), True except Exception: return text, False return text, False class PromptLoader: def __init__(self): self.config_file = "prompts_config.json" self.data = self.load_data() def load_data(self): try: path = hf_hub_download(repo_id=DATASET_ID, filename=self.config_file, repo_type="dataset", token=HF_TOKEN, local_files_only=True) return json.load(open(path, "r")) except: try: path = hf_hub_download(repo_id=DATASET_ID, filename=self.config_file, repo_type="dataset", token=HF_TOKEN) return json.load(open(path, "r")) except: return { "subjects": ["cyberpunk city", "magical forest", "robot cat"], "environments": ["in rain", "under moonlight", "in space"], "lighting": ["neon", "cinematic", "natural"] } def get_random_prompt(self): s = random.choice(self.data.get("subjects", ["cat"])) e = random.choice(self.data.get("environments", ["home"])) l = random.choice(self.data.get("lighting", ["day"])) return f"Illustration of {s} {e}, {l}." class LocalUpscaler: def __init__(self): self.model_path = "FSRCNN_x2.pb" def process(self, image_bytes): if OPENCV_AVAILABLE and not os.path.exists(self.model_path): try: r = requests.get("https://github.com/Saafke/FSRCNN_Tensorflow/raw/master/models/FSRCNN_x2.pb", timeout=30) if r.status_code == 200: with open(self.model_path, "wb") as f: f.write(r.content) except: pass if OPENCV_AVAILABLE and os.path.exists(self.model_path): try: nparr = np.frombuffer(image_bytes, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) sr = cv2.dnn_superres.DnnSuperResImpl_create() sr.readModel(self.model_path) sr.setModel("fsrcnn", 2) result = sr.upsample(img) is_success, buffer = cv2.imencode(".png", result) if is_success: return buffer.tobytes(), "AI FSRCNN" except: pass try: with Image.open(io.BytesIO(image_bytes)) as img: upscaled = img.resize((img.width * 2, img.height * 2), Image.Resampling.LANCZOS) buf = io.BytesIO() upscaled.save(buf, format="PNG") return buf.getvalue(), "HQ Lanczos" except Exception as e: raise ValueError(f"Upscaling failed: {e}") class SocialPosterGenerator: def __init__(self): self.font_path = "SimHei.ttf" self._ensure_font() def _ensure_font(self): if not os.path.exists(self.font_path): try: r = requests.get("https://github.com/StellarCN/scp_zh/raw/master/fonts/SimHei.ttf", timeout=30) with open(self.font_path, "wb") as f: f.write(r.content) except: pass def _get_text_width(self, text, font, draw): return draw.textlength(text, font=font) def create_card(self, image_bytes, prompt_text, qr_content, footer_text="AI 灵感绘图 PRO"): try: base_img = Image.open(io.BytesIO(image_bytes)).convert("RGBA") except: return None padding = 50 text_area_height = 340 canvas_width = base_img.width canvas_height = base_img.height + text_area_height poster = Image.new("RGBA", (canvas_width, canvas_height), (255, 255, 255, 255)) poster.paste(base_img, (0, 0)) draw = ImageDraw.Draw(poster) try: font_title = ImageFont.truetype(self.font_path, 36) font_text = ImageFont.truetype(self.font_path, 24) font_footer = ImageFont.truetype(self.font_path, 20) except: font_title = font_text = font_footer = ImageFont.load_default() draw.text((padding, base_img.height + 40), "🎨 AI 创意画作", font=font_title, fill="#222222") qr_size = 150 try: qr = qrcode.QRCode(box_size=4, border=1) qr.add_data(qr_content or "https://huggingface.co/spaces") qr.make(fit=True) qr_img = qr.make_image(fill_color="black", back_color="white").convert("RGBA").resize((qr_size, qr_size)) poster.paste(qr_img, (canvas_width - qr_size - padding, base_img.height + (text_area_height - qr_size) // 2)) except: pass text_width_limit = canvas_width - padding * 2 - qr_size - 20 lines = [] current_line = "" clean_text = prompt_text.replace('\n', ' ').replace('\r', '') for char in clean_text: if self._get_text_width(current_line + char, font_text, draw) <= text_width_limit: current_line += char else: lines.append(current_line) current_line = char if current_line: lines.append(current_line) if len(lines) > 4: lines = lines[:4]; lines[-1] = lines[-1][:-3] + "..." current_h = base_img.height + 100 for line in lines: draw.text((padding, current_h), line, font=font_text, fill="#555555") current_h += 36 date_str = datetime.datetime.now().strftime('%Y-%m-%d') draw.text((padding, canvas_height - 50), f"{footer_text} · {date_str}", font=font_footer, fill="#999999") buf = io.BytesIO() poster.convert("RGB").save(buf, format="JPEG", quality=95) # GC helper base_img.close() poster.close() return buf.getvalue()