| 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
|
|
|
|
|
| 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:
|
|
|
| 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)
|
|
|
| base_img.close()
|
| poster.close()
|
| return buf.getvalue() |
