1242 lines
50 KiB
Python
1242 lines
50 KiB
Python
import os
|
||
import sys
|
||
import json
|
||
import re
|
||
import requests
|
||
import threading
|
||
import time
|
||
import base64
|
||
import io
|
||
import webview
|
||
import random
|
||
import subprocess
|
||
import concurrent.futures
|
||
from PIL import Image, PngImagePlugin
|
||
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
|
||
import torch
|
||
import replicate
|
||
|
||
# ------------- Configuration -------------
|
||
|
||
OLLAMA_URL = "http://localhost:11434/api/generate"
|
||
MODEL = "mistral:latest" # or "mistral-small3.1:24b"
|
||
|
||
# Path to your diffusers-converted model
|
||
MODEL_PATH = "/Volumes/SD/ML-Models/diffusers/dreamshaper_8_diffusers"
|
||
|
||
# Where to save generated images
|
||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
# ────────────────────────────────────────────────────────────────────────────
|
||
# Store everything under web/output instead of project_root/output
|
||
WEB_DIR = os.path.join(BASE_DIR, "web")
|
||
OUTPUT_DIR = os.path.join(WEB_DIR, "output")
|
||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||
APP_SUPPORT_DIR = os.path.join(
|
||
os.path.expanduser("~"),
|
||
"Library",
|
||
"Application Support",
|
||
"3d-model-generator",
|
||
)
|
||
SETTINGS_PATH = os.path.join(APP_SUPPORT_DIR, "settings.json")
|
||
# ────────────────────────────────────────────────────────────────────────────
|
||
|
||
IMAGE_BACKEND_LOCAL = "local"
|
||
IMAGE_BACKEND_REPLICATE = "replicate"
|
||
DEFAULT_IMAGE_BACKEND = IMAGE_BACKEND_LOCAL
|
||
DEFAULT_REPLICATE_IMAGE_MODEL = "google/imagen-4-fast"
|
||
REPLICATE_IMAGE_TIMEOUT = 10 * 60
|
||
REPLICATE_IMAGE_POLL_INTERVAL = 2.0
|
||
|
||
REPLICATE_IMAGE_MODELS = {
|
||
"black-forest-labs/flux-schnell": {
|
||
"label": "FLUX.1 schnell",
|
||
"price": "$3 / 1000 output images (~$0.003/image)",
|
||
"aspect_ratios": ["1:1", "16:9", "21:9", "3:2", "2:3", "4:5", "5:4", "3:4", "4:3", "9:16", "9:21"],
|
||
},
|
||
"google/imagen-4-fast": {
|
||
"label": "Google Imagen 4 Fast",
|
||
"price": "$0.02 / output image",
|
||
"aspect_ratios": ["1:1", "9:16", "16:9", "3:4", "4:3"],
|
||
},
|
||
"bytedance/seedream-4": {
|
||
"label": "ByteDance Seedream 4",
|
||
"price": "$0.03 / output image",
|
||
"aspect_ratios": ["1:1", "4:3", "3:4", "16:9", "9:16", "3:2", "2:3", "21:9"],
|
||
},
|
||
}
|
||
|
||
|
||
def load_app_settings() -> dict:
|
||
try:
|
||
with open(SETTINGS_PATH, "r", encoding="utf-8") as f:
|
||
settings = json.load(f)
|
||
return settings if isinstance(settings, dict) else {}
|
||
except FileNotFoundError:
|
||
return {}
|
||
except Exception:
|
||
return {}
|
||
|
||
|
||
def save_app_settings(settings: dict) -> None:
|
||
os.makedirs(APP_SUPPORT_DIR, exist_ok=True)
|
||
with open(SETTINGS_PATH, "w", encoding="utf-8") as f:
|
||
json.dump(settings, f, indent=2)
|
||
|
||
# ------------- Prompt Templates -------------
|
||
|
||
META_PROMPT = """
|
||
You are an expert at writing concise, detailed Stable Diffusion prompts for 3D rendered objects, as seen in professional game development.
|
||
|
||
When I give you an object name (which may already contain style or material cues, e.g. "low poly farmer" or "pixel art sword"), follow these rules:
|
||
|
||
1. Parse the object name:
|
||
- If the object name includes a style, technique, or specific look (e.g., low poly, pixel art, voxel, cartoon, realistic, hand painted, clay, wireframe, etc.), make that style the focus of the prompt.
|
||
- Do NOT add conflicting style keywords. (E.g. do NOT add high poly if the input is low poly.)
|
||
- If the object name is generic (like farmer or spaceship), you may suggest professional, detailed, high-quality, realistic or stylized game art (choose one style, do NOT mix).
|
||
- Add 2–3 relevant visual or material details to the object, based on typical game art conventions.
|
||
|
||
2. Medium & Style:
|
||
- If the object name includes a medium, engine, or tool (Unreal, Octane, Unity, voxel, etc.), use it and do NOT add others.
|
||
- Otherwise, you may choose (but do NOT mix) one or two: concept art, 3D render, game asset, professional game designer, digital sculpture, etc.
|
||
|
||
3. Presentation:
|
||
- State that the object is fully visible, not cropped, and centered from a neutral angle.
|
||
- No background: isolated on pure white or transparent background.
|
||
|
||
4. Lighting & Quality:
|
||
- Only add studio lighting, sharp focus, clean silhouette, ultra detailed, 8K, etc. if compatible with the style.
|
||
- Avoid depth of field, blur, or effects that obscure details, unless specifically requested.
|
||
|
||
5. Negative Prompt:
|
||
- Always add: Negative prompt: blurry, lowres, bad anatomy, distorted proportions, background, scenery, environment, artifacting, watermark, text, cropped, partial view, depth of field, out of frame, blur, vignette
|
||
|
||
Output Format:
|
||
- Line 1: All positive prompt keywords (subject, details, medium, style, lighting, quality, presentation, background).
|
||
- Line 2: Negative prompt: then the negative keywords above.
|
||
|
||
Instructions:
|
||
- If the input object already specifies a style or technique, your prompt must reflect that and avoid all contradictory terms.
|
||
- Always keep the style consistent.
|
||
- Avoid unnecessary repetition.
|
||
- Never add stylistic terms that contradict or dilute the user’s intention.
|
||
- The positive prompt must not contain more than 300 characters. Make it just fit 300 characters, not too long, not too short but perfect for a Stable Diffusion prompt.
|
||
- If the prompt is short, come up with creative ideas on how to make the character or object more intricate and interesting, suitable for game design. Add creative details to what the presented things might look like!
|
||
- Use positive descriptions of what should be visible in the positive prompt, don't use negations there.
|
||
- Most importantly, the prompts must make sure in some way that the object or character is FULLY visible (not cropped) on a blank background.
|
||
|
||
Now, generate a Stable Diffusion prompt for the object: {{OBJECT}}
|
||
"""
|
||
|
||
SECOND_PROMPT_TEMPLATE = """
|
||
You are a prompt-to-JSON converter for image generation tasks.
|
||
|
||
Given a Stable Diffusion prompt, extract the positive and negative prompts.
|
||
If the prompt is longer than 300 characters, summarize it so that it's less than 300 characters long.
|
||
If you need to shorten your prompt, always remove less relevant visual details first.
|
||
Never remove the presentation instructions (e.g., fully visible, isolated, centered, white background); these must always remain at the end of the positive prompt.
|
||
Remove words that are followed by a semicolon (i.e. "Presentation:" or "Lighting & Quality:") - these things have no business in a Stable Diffusion prompt.
|
||
|
||
Also, analyze the subject:
|
||
- If it is a single character, an upright (bipedal) creature, full object, a person, or anything similar, set "dimensions" to [512,768] (vertical). (In that case of a bipedal creature / human, you might also add "A-Pose" to the prompt!)
|
||
- If it is a landscape, wide object, multi-character group, or scene, or a creature that isn't tall (four-legged creatures for example), set "dimensions" to [768,512] (horizontal).
|
||
- If it is something square or best seen as a square (e.g., shield, logo, face, emblem, single centered item or a box or round-shaped creature), set "dimensions" to [768,768] (square).
|
||
|
||
Return a JSON object with these fields:
|
||
{
|
||
"positive prompt": "...",
|
||
"negative prompt": "...",
|
||
"dimensions": [W,H]
|
||
}
|
||
|
||
Respond only with the JSON.
|
||
Here is the Stable Diffusion prompt:
|
||
|
||
{PROMPT}
|
||
"""
|
||
|
||
# ------------- Helper Functions -------------
|
||
|
||
def stream_ollama(prompt: str):
|
||
"""
|
||
Stream tokens from Ollama.
|
||
Yields the cumulative response string after each new chunk arrives.
|
||
"""
|
||
payload = {"model": MODEL, "prompt": prompt, "stream": True}
|
||
with requests.post(OLLAMA_URL, json=payload, stream=True, timeout=600) as r:
|
||
r.raise_for_status()
|
||
full_response = ""
|
||
for line in r.iter_lines():
|
||
if line:
|
||
try:
|
||
chunk = json.loads(line.decode("utf-8"))["response"]
|
||
except Exception:
|
||
chunk = line.decode("utf-8")
|
||
full_response += chunk
|
||
yield full_response
|
||
yield full_response.strip()
|
||
|
||
def get_json_from_llm(prompt: str):
|
||
"""
|
||
Stream the JSON conversion call to Ollama (no LangChain).
|
||
Yields the cumulative response string after each chunk arrives.
|
||
"""
|
||
payload = {"model": MODEL, "prompt": prompt, "stream": True}
|
||
with requests.post(OLLAMA_URL, json=payload, stream=True, timeout=600) as r:
|
||
r.raise_for_status()
|
||
full_response = ""
|
||
for line in r.iter_lines():
|
||
if line:
|
||
try:
|
||
chunk = json.loads(line.decode("utf-8"))["response"]
|
||
except Exception:
|
||
chunk = line.decode("utf-8")
|
||
full_response += chunk
|
||
yield full_response
|
||
yield full_response.strip()
|
||
|
||
def sanitize_filename(s: str) -> str:
|
||
s = s.strip().lower().replace(" ", "_")
|
||
return re.sub(r'[^a-z0-9_]+', '', s)
|
||
|
||
def get_next_available_filename(base: str, ext: str = ".png") -> str:
|
||
i = 1
|
||
while True:
|
||
fname = f"{base}-{i}{ext}"
|
||
# → hier steht noch fpath = f"output/{fname}"
|
||
fpath = os.path.join(OUTPUT_DIR, fname)
|
||
if not os.path.exists(fpath):
|
||
return fpath
|
||
i += 1
|
||
|
||
def image_to_base64(path: str) -> str:
|
||
"""
|
||
Read an image from disk and return a data: URI string.
|
||
"""
|
||
with open(path, "rb") as f:
|
||
b64 = base64.b64encode(f.read()).decode("utf-8")
|
||
return f"data:image/png;base64,{b64}"
|
||
|
||
|
||
def normalize_image_backend(value: str | None) -> str:
|
||
value = (value or DEFAULT_IMAGE_BACKEND).strip().lower()
|
||
if value in {IMAGE_BACKEND_LOCAL, IMAGE_BACKEND_REPLICATE}:
|
||
return value
|
||
return DEFAULT_IMAGE_BACKEND
|
||
|
||
|
||
def normalize_replicate_image_model(value: str | None) -> str:
|
||
value = (value or DEFAULT_REPLICATE_IMAGE_MODEL).strip()
|
||
if value in REPLICATE_IMAGE_MODELS:
|
||
return value
|
||
return DEFAULT_REPLICATE_IMAGE_MODEL
|
||
|
||
|
||
def nearest_aspect_ratio(width: int, height: int, allowed: list[str]) -> str:
|
||
if width <= 0 or height <= 0:
|
||
return "1:1" if "1:1" in allowed else allowed[0]
|
||
|
||
target = width / height
|
||
|
||
def ratio_value(value: str) -> float:
|
||
left, right = value.split(":", 1)
|
||
return float(left) / float(right)
|
||
|
||
return min(allowed, key=lambda option: abs(ratio_value(option) - target))
|
||
|
||
|
||
def compose_replicate_prompt(prompt_json: dict) -> str:
|
||
positive = str(prompt_json.get("positive prompt", "")).strip()
|
||
negative = str(prompt_json.get("negative prompt", "")).strip()
|
||
if not negative:
|
||
return positive
|
||
return f"{positive}\nAvoid: {negative}"
|
||
|
||
|
||
def extract_replicate_output_urls(output) -> list[str]:
|
||
if output is None:
|
||
return []
|
||
if hasattr(output, "url"):
|
||
return [str(output.url)]
|
||
if isinstance(output, str):
|
||
return [output]
|
||
if isinstance(output, (list, tuple)):
|
||
urls = []
|
||
for item in output:
|
||
urls.extend(extract_replicate_output_urls(item))
|
||
return urls
|
||
if isinstance(output, dict):
|
||
urls = []
|
||
for item in output.values():
|
||
urls.extend(extract_replicate_output_urls(item))
|
||
return urls
|
||
return []
|
||
|
||
|
||
def load_image_from_url(url: str) -> Image.Image:
|
||
if url.startswith("data:"):
|
||
_, encoded = url.split(",", 1)
|
||
data = base64.b64decode(encoded)
|
||
else:
|
||
response = requests.get(url, timeout=REPLICATE_IMAGE_TIMEOUT)
|
||
response.raise_for_status()
|
||
data = response.content
|
||
|
||
img = Image.open(io.BytesIO(data))
|
||
if img.mode in {"RGBA", "LA"} or (img.mode == "P" and "transparency" in img.info):
|
||
return img.convert("RGBA")
|
||
return img.convert("RGB")
|
||
|
||
|
||
def save_generated_image(img: Image.Image, prompt_json: dict, object_name: str, idx: int, batch_count: int) -> str:
|
||
fname = get_next_available_filename(f"{sanitize_filename(object_name)}-{idx}")
|
||
metadata = dict(prompt_json, created=int(time.time()), batch_count=batch_count)
|
||
meta = PngImagePlugin.PngInfo()
|
||
meta.add_text("sd_json", json.dumps(metadata))
|
||
img.save(fname, pnginfo=meta)
|
||
with open(fname + ".json", "w", encoding="utf-8") as f:
|
||
json.dump(metadata, f, indent=2)
|
||
return fname
|
||
|
||
|
||
def call_stable_diffusion(prompt_json: dict, user_input: str) -> list[str]:
|
||
"""
|
||
Generate 'n_imgs' images from diffusers, save to disk with embedded prompt JSON,
|
||
write sidecar .json, and return the file paths.
|
||
"""
|
||
if torch.backends.mps.is_available():
|
||
device = "mps"
|
||
else:
|
||
device = "cpu"
|
||
|
||
# Load pipeline
|
||
pipe = StableDiffusionPipeline.from_pretrained(
|
||
MODEL_PATH,
|
||
torch_dtype=torch.float32,
|
||
safety_checker=None,
|
||
local_files_only=True
|
||
)
|
||
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
|
||
pipe = pipe.to(device)
|
||
pipe.set_progress_bar_config(disable=True)
|
||
pipe.enable_attention_slicing()
|
||
|
||
pos_prompt = prompt_json["positive prompt"]
|
||
neg_prompt = prompt_json["negative prompt"]
|
||
width, height = prompt_json["dimensions"]
|
||
steps = prompt_json.get("steps", 30)
|
||
guidance_scale = prompt_json.get("cfg_scale", 6.5)
|
||
batch_count = int(prompt_json.get("batch_count", 4)) # Standard: 4
|
||
|
||
images = pipe(
|
||
prompt=[pos_prompt] * batch_count,
|
||
negative_prompt=[neg_prompt] * batch_count,
|
||
width=width,
|
||
height=height,
|
||
num_inference_steps=steps,
|
||
guidance_scale=guidance_scale
|
||
).images
|
||
|
||
img_base = sanitize_filename(user_input)
|
||
out_files = []
|
||
import time
|
||
for idx, img in enumerate(images, start=1):
|
||
filename = get_next_available_filename(f"{img_base}-{idx}")
|
||
|
||
# Timestamp und batch_count speichern
|
||
prompt_json_with_time = dict(prompt_json)
|
||
prompt_json_with_time["created"] = int(time.time())
|
||
prompt_json_with_time["batch_count"] = batch_count
|
||
|
||
# Metadata und Sidecar wie gehabt
|
||
meta = PngImagePlugin.PngInfo()
|
||
meta.add_text("sd_json", json.dumps(prompt_json_with_time))
|
||
img.save(filename, pnginfo=meta)
|
||
sidecar_path = filename + ".json"
|
||
with open(sidecar_path, "w", encoding="utf-8") as f:
|
||
json.dump(prompt_json_with_time, f, indent=2)
|
||
out_files.append(filename)
|
||
return out_files
|
||
|
||
|
||
# ------------- JSON “Fuzzy” Validator -------------
|
||
|
||
def validate_and_fix_json(raw: str) -> tuple[dict, str] | tuple[None, None]:
|
||
"""
|
||
1) Try a strict json.loads(raw).
|
||
2) If that fails, attempt to fix:
|
||
- Remove trailing commas (e.g. {"a":1,} -> {"a":1})
|
||
- Replace single quotes with double quotes (naïve)
|
||
3) If the “fixed” string parses, return (parsed_dict, fixed_string).
|
||
Otherwise, return (None, None).
|
||
"""
|
||
try:
|
||
parsed = json.loads(raw)
|
||
return parsed, raw
|
||
except Exception:
|
||
# Attempt quick fixes:
|
||
s = raw
|
||
# 1) Remove any trailing commas before } or ]
|
||
s = re.sub(r",\s*([}\]])", r"\1", s)
|
||
# 2) Replace single quotes around keys/values with double quotes (naïve)
|
||
s = re.sub(r"(?P<pre>[:\s])'(?P<val>[^']*)'", r'\1"\g<val>"', s)
|
||
try:
|
||
parsed = json.loads(s)
|
||
return parsed, s
|
||
except Exception:
|
||
return None, None
|
||
|
||
# ------------- PyWebview API -------------
|
||
|
||
class Api:
|
||
def __init__(self):
|
||
self.window = None
|
||
self.last_object = ""
|
||
self.last_prompt_json = None
|
||
self._abort_sd = False
|
||
self._placeholders_ready_events = {}
|
||
|
||
# Keep track of file modification times
|
||
self._mtimes: dict[str, float] = {}
|
||
# Start folder watcher thread
|
||
watcher = threading.Thread(target=self._monitor_folder, daemon=True)
|
||
watcher.start()
|
||
|
||
def set_window(self, win):
|
||
self.window = win
|
||
|
||
def generate_prompt(self, object_name: str):
|
||
"""
|
||
Runs only the LLM -> JSON steps (no Stable Diffusion).
|
||
"""
|
||
self.last_object = object_name
|
||
self.last_prompt_json = None
|
||
threading.Thread(target=self._run_llm_to_json, args=(object_name,), daemon=True).start()
|
||
return {"status": "started_prompt"}
|
||
|
||
def _run_llm_to_json(self, object_name: str):
|
||
# 1) LLM streaming
|
||
llm_text = ""
|
||
llm_prompt = META_PROMPT.replace("{{OBJECT}}", object_name)
|
||
for chunk in stream_ollama(llm_prompt):
|
||
llm_text = chunk
|
||
self._js("window.on_llm_output", llm_text)
|
||
self._js("window.on_llm_done")
|
||
|
||
# 2) JSON streaming
|
||
json_text = ""
|
||
json_prompt = SECOND_PROMPT_TEMPLATE.replace("{PROMPT}", llm_text)
|
||
for chunk in get_json_from_llm(json_prompt):
|
||
json_text = chunk
|
||
self._js("window.on_json_output", json_text)
|
||
self._js("window.on_json_done")
|
||
|
||
# 3) Parse JSON
|
||
try:
|
||
prompt_dict = json.loads(json_text)
|
||
except Exception as e:
|
||
self._js("window.on_error", f"JSON parse error: {e}")
|
||
return
|
||
|
||
# 4) Augment with SD settings
|
||
prompt_dict["checkpoint"] = "RealismPlus/dreamshaper_8.safetensors"
|
||
prompt_dict["vae"] = "Automatic"
|
||
prompt_dict["sampler"] = "DPM++ 2M Karras"
|
||
prompt_dict["steps"] = 30
|
||
prompt_dict["cfg_scale"] = 6.5
|
||
prompt_dict["batch_count"] = 4
|
||
|
||
# 5) Send full JSON
|
||
full_json_str = json.dumps(prompt_dict, indent=2)
|
||
self._js("window.on_json_output", full_json_str)
|
||
|
||
# 6) Store for images
|
||
self.last_prompt_json = prompt_dict
|
||
|
||
def generate_images(self, json_str: str):
|
||
"""
|
||
Called from JS when “Generate Images” is clicked.
|
||
1) Validate/fix JSON (fuzzy).
|
||
2) Ensure “positive prompt” is nonempty.
|
||
3) Auto fill missing keys (dimensions, checkpoint, etc.)
|
||
4) Kick off Stable Diffusion.
|
||
"""
|
||
parsed, fixed = validate_and_fix_json(json_str)
|
||
if parsed is None:
|
||
self._js("window.on_error", "Cannot parse JSON (even after auto-fix).")
|
||
return {"status": "invalid_json"}
|
||
|
||
pos = parsed.get("positive prompt", "").strip()
|
||
if not pos:
|
||
self._js("window.on_error", "“positive prompt” must not be empty.")
|
||
return {"status": "empty_positive"}
|
||
|
||
# Auto-fill missing/malformed keys
|
||
dims = parsed.get("dimensions")
|
||
if not (
|
||
isinstance(dims, list)
|
||
and len(dims) == 2
|
||
and all(isinstance(x, int) for x in dims)
|
||
):
|
||
parsed["dimensions"] = [768, 768]
|
||
parsed.setdefault("checkpoint", "RealismPlus/dreamshaper_8.safetensors")
|
||
parsed.setdefault("vae", "Automatic")
|
||
parsed.setdefault("sampler", "DPM++ 2M Karras")
|
||
parsed.setdefault("steps", 30)
|
||
parsed.setdefault("cfg_scale", 6.5)
|
||
parsed.setdefault("batch_count", 4)
|
||
|
||
# If we “fixed” it, push corrected JSON back
|
||
full_json_str = json.dumps(parsed, indent=2)
|
||
if fixed is None or full_json_str != json_str:
|
||
self._js("window.on_json_output", full_json_str)
|
||
|
||
# Store and infer last_object if not yet set
|
||
self.last_prompt_json = parsed
|
||
if not self.last_object:
|
||
# Derive a base name from the positive prompt (before first comma)
|
||
cand = pos.split(",")[0]
|
||
base = re.sub(r"\s+", "_", cand.strip().lower())
|
||
base = re.sub(r"[^a-z0-9_]+", "", base)
|
||
self.last_object = base or "image"
|
||
|
||
# Launch SD thread – vorher Abort-Flag zurücksetzen
|
||
self._abort_sd = False
|
||
threading.Thread(target=self._run_sd_only, daemon=True).start()
|
||
return {"status": "started_images"}
|
||
|
||
def _get_image_generation_backend(self) -> str:
|
||
return normalize_image_backend(load_app_settings().get("image_generation_backend"))
|
||
|
||
def _get_replicate_image_model(self) -> str:
|
||
return normalize_replicate_image_model(load_app_settings().get("replicate_image_model"))
|
||
|
||
def _run_local_diffusion_images(self, batch_count: int) -> list[Image.Image]:
|
||
device = "mps" if torch.backends.mps.is_available() else "cpu"
|
||
|
||
pipe = StableDiffusionPipeline.from_pretrained(
|
||
MODEL_PATH, torch_dtype=torch.float32,
|
||
safety_checker=None, local_files_only=True
|
||
)
|
||
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
|
||
pipe = pipe.to(device)
|
||
pipe.set_progress_bar_config(disable=True)
|
||
pipe.enable_attention_slicing()
|
||
|
||
pos = self.last_prompt_json["positive prompt"]
|
||
neg = self.last_prompt_json["negative prompt"]
|
||
w, h = self.last_prompt_json["dimensions"]
|
||
steps = self.last_prompt_json.get("steps", 30)
|
||
cfg = self.last_prompt_json.get("cfg_scale", 6.5)
|
||
|
||
def _check_abort(*args, **kwargs):
|
||
if self._abort_sd:
|
||
raise RuntimeError("User aborted")
|
||
return {}
|
||
|
||
outputs = pipe(
|
||
prompt=[pos] * batch_count,
|
||
negative_prompt=[neg] * batch_count,
|
||
width=w, height=h,
|
||
num_inference_steps=steps,
|
||
guidance_scale=cfg,
|
||
callback_on_step_end=_check_abort,
|
||
)
|
||
return outputs.images
|
||
|
||
def _build_replicate_image_input(self, model_id: str, num_outputs: int) -> dict:
|
||
model = REPLICATE_IMAGE_MODELS[model_id]
|
||
w, h = self.last_prompt_json["dimensions"]
|
||
ratio = nearest_aspect_ratio(w, h, model["aspect_ratios"])
|
||
prompt = compose_replicate_prompt(self.last_prompt_json)
|
||
|
||
if model_id == "black-forest-labs/flux-schnell":
|
||
steps = int(self.last_prompt_json.get("steps", 4))
|
||
return {
|
||
"prompt": prompt,
|
||
"aspect_ratio": ratio,
|
||
"num_outputs": max(1, min(int(num_outputs), 4)),
|
||
"num_inference_steps": max(1, min(steps, 4)),
|
||
"output_format": "png",
|
||
"go_fast": True,
|
||
"megapixels": "1",
|
||
}
|
||
|
||
if model_id == "google/imagen-4-fast":
|
||
return {
|
||
"prompt": prompt,
|
||
"aspect_ratio": ratio,
|
||
"output_format": "png",
|
||
"safety_filter_level": "block_only_high",
|
||
}
|
||
|
||
if model_id == "bytedance/seedream-4":
|
||
return {
|
||
"prompt": prompt,
|
||
"aspect_ratio": ratio,
|
||
"size": "2K",
|
||
"max_images": 1,
|
||
"enhance_prompt": False,
|
||
"sequential_image_generation": "disabled",
|
||
}
|
||
|
||
raise RuntimeError(f"Unsupported Replicate image model: {model_id}")
|
||
|
||
def _run_replicate_prediction(self, model_id: str, input_payload: dict):
|
||
token = self._get_replicate_api_token()
|
||
if not token:
|
||
raise RuntimeError("Replicate API token missing. Open settings and add it.")
|
||
|
||
client = replicate.Client(api_token=token)
|
||
client.poll_interval = REPLICATE_IMAGE_POLL_INTERVAL
|
||
prediction = client.models.predictions.create(
|
||
model=model_id,
|
||
input=input_payload,
|
||
wait=False,
|
||
)
|
||
print(f"[Replicate Image] Prediction started: {prediction.id} ({prediction.status})")
|
||
|
||
deadline = time.monotonic() + REPLICATE_IMAGE_TIMEOUT
|
||
last_status = prediction.status
|
||
while prediction.status not in {"succeeded", "failed", "canceled"}:
|
||
if self._abort_sd:
|
||
try:
|
||
prediction.cancel()
|
||
except Exception:
|
||
pass
|
||
raise RuntimeError("User aborted")
|
||
|
||
if time.monotonic() >= deadline:
|
||
raise TimeoutError(
|
||
f"Timed out waiting for Replicate image prediction {prediction.id} "
|
||
f"after {REPLICATE_IMAGE_TIMEOUT // 60} minutes."
|
||
)
|
||
|
||
time.sleep(client.poll_interval)
|
||
prediction.reload()
|
||
if prediction.status != last_status:
|
||
print(f"[Replicate Image] {prediction.id}: {prediction.status}")
|
||
last_status = prediction.status
|
||
|
||
if prediction.status != "succeeded":
|
||
detail = prediction.error or prediction.logs or f"status={prediction.status}"
|
||
raise RuntimeError(f"Replicate image prediction {prediction.id} failed: {detail}")
|
||
|
||
return prediction.output
|
||
|
||
def _run_replicate_images(self, batch_count: int) -> list[Image.Image]:
|
||
model_id = self._get_replicate_image_model()
|
||
images: list[Image.Image] = []
|
||
|
||
if model_id == "black-forest-labs/flux-schnell":
|
||
remaining = batch_count
|
||
while remaining > 0:
|
||
if self._abort_sd:
|
||
raise RuntimeError("User aborted")
|
||
chunk = min(remaining, 4)
|
||
output = self._run_replicate_prediction(
|
||
model_id,
|
||
self._build_replicate_image_input(model_id, chunk),
|
||
)
|
||
urls = extract_replicate_output_urls(output)
|
||
if len(urls) < chunk:
|
||
raise RuntimeError(f"Replicate returned {len(urls)} image(s), expected {chunk}.")
|
||
images.extend(load_image_from_url(url) for url in urls[:chunk])
|
||
remaining -= chunk
|
||
return images
|
||
|
||
for _ in range(batch_count):
|
||
if self._abort_sd:
|
||
raise RuntimeError("User aborted")
|
||
output = self._run_replicate_prediction(
|
||
model_id,
|
||
self._build_replicate_image_input(model_id, 1),
|
||
)
|
||
urls = extract_replicate_output_urls(output)
|
||
if not urls:
|
||
raise RuntimeError("Replicate did not return an image URL.")
|
||
images.append(load_image_from_url(urls[0]))
|
||
|
||
return images
|
||
|
||
def _run_sd_only(self):
|
||
if not self.last_prompt_json:
|
||
self._js("window.on_error", "No valid JSON prompt available.")
|
||
return
|
||
|
||
batch_id = random.randint(1_000_000, 9_999_999)
|
||
self._last_batch_id = batch_id
|
||
batch_count = int(self.last_prompt_json.get("batch_count", 4))
|
||
|
||
# 1) Bail out if user already hit “Stop”
|
||
if self._abort_sd:
|
||
return
|
||
|
||
# 2) Show placeholders in the UI
|
||
self._js("window.show_placeholders", batch_count, batch_id)
|
||
event = threading.Event()
|
||
self._placeholders_ready_events[batch_id] = event
|
||
|
||
# 3) Wait until JS signals placeholders are in place
|
||
while not event.is_set():
|
||
time.sleep(0.02)
|
||
if self._abort_sd:
|
||
self._js("window.remove_placeholders", batch_id)
|
||
return
|
||
|
||
# 4) Set up and run the selected image generator.
|
||
try:
|
||
if self._get_image_generation_backend() == IMAGE_BACKEND_REPLICATE:
|
||
images = self._run_replicate_images(batch_count)
|
||
else:
|
||
images = self._run_local_diffusion_images(batch_count)
|
||
|
||
except RuntimeError as e:
|
||
# user‐initiated abort
|
||
if str(e) == "User aborted":
|
||
self._js("window.on_generation_aborted")
|
||
self._js("window.remove_placeholders", batch_id)
|
||
return
|
||
# or some other error
|
||
self._js("window.remove_placeholders", batch_id)
|
||
self._js("window.on_error", f"Image generation error: {e}")
|
||
return
|
||
except Exception as e:
|
||
self._js("window.remove_placeholders", batch_id)
|
||
self._js("window.on_error", f"Image generation error: {e}")
|
||
return
|
||
|
||
# 5) Replace placeholders one by one
|
||
for idx, img in enumerate(images):
|
||
if self._abort_sd:
|
||
self._js("window.remove_placeholders", batch_id)
|
||
return
|
||
|
||
fname = save_generated_image(
|
||
img,
|
||
self.last_prompt_json,
|
||
self.last_object,
|
||
idx + 1,
|
||
batch_count,
|
||
)
|
||
|
||
rel = os.path.relpath(fname, WEB_DIR)
|
||
self._js("window.replace_placeholder", idx, rel, batch_id)
|
||
|
||
# 6) Tear down
|
||
self._js("window.remove_placeholders", batch_id)
|
||
self._js("window.on_image_gen_done")
|
||
del self._placeholders_ready_events[batch_id]
|
||
|
||
def stop_generation(self):
|
||
self._abort_sd = True
|
||
# Do NOT remove placeholders here—wait until the abort is processed
|
||
return {"status": "aborting"}
|
||
|
||
def placeholders_ready(self, batch_id):
|
||
if batch_id in self._placeholders_ready_events:
|
||
self._placeholders_ready_events[batch_id].set()
|
||
return {"status": "acknowledged"}
|
||
|
||
def get_settings(self):
|
||
settings = load_app_settings()
|
||
settings_token = str(settings.get("replicate_api_token", "")).strip()
|
||
env_token = os.environ.get("REPLICATE_API_TOKEN", "").strip()
|
||
source = "settings" if settings_token else "environment" if env_token else ""
|
||
return {
|
||
"has_replicate_api_token": bool(settings_token or env_token),
|
||
"replicate_api_token_source": source,
|
||
"replicate_api_token": "",
|
||
"image_generation_backend": normalize_image_backend(settings.get("image_generation_backend")),
|
||
"replicate_image_model": normalize_replicate_image_model(settings.get("replicate_image_model")),
|
||
"replicate_image_models": [
|
||
{
|
||
"id": model_id,
|
||
"label": model["label"],
|
||
"price": model["price"],
|
||
}
|
||
for model_id, model in REPLICATE_IMAGE_MODELS.items()
|
||
],
|
||
}
|
||
|
||
def save_settings(self, settings: dict):
|
||
current = load_app_settings()
|
||
token = str(settings.get("replicate_api_token", "")).strip()
|
||
if token:
|
||
current["replicate_api_token"] = token
|
||
|
||
current["image_generation_backend"] = normalize_image_backend(
|
||
settings.get("image_generation_backend")
|
||
)
|
||
current["replicate_image_model"] = normalize_replicate_image_model(
|
||
settings.get("replicate_image_model")
|
||
)
|
||
save_app_settings(current)
|
||
return {
|
||
"status": "saved",
|
||
"has_replicate_api_token": bool(self._get_replicate_api_token()),
|
||
"image_generation_backend": current["image_generation_backend"],
|
||
"replicate_image_model": current["replicate_image_model"],
|
||
}
|
||
|
||
def clear_replicate_api_token(self):
|
||
current = load_app_settings()
|
||
current.pop("replicate_api_token", None)
|
||
save_app_settings(current)
|
||
env_token = os.environ.get("REPLICATE_API_TOKEN", "").strip()
|
||
return {
|
||
"status": "cleared",
|
||
"has_replicate_api_token": bool(env_token),
|
||
"replicate_api_token_source": "environment" if env_token else "",
|
||
"image_generation_backend": normalize_image_backend(current.get("image_generation_backend")),
|
||
"replicate_image_model": normalize_replicate_image_model(current.get("replicate_image_model")),
|
||
}
|
||
|
||
def _get_replicate_api_token(self):
|
||
settings_token = str(load_app_settings().get("replicate_api_token", "")).strip()
|
||
return settings_token or os.environ.get("REPLICATE_API_TOKEN", "").strip()
|
||
|
||
def _subprocess_env(self):
|
||
env = os.environ.copy()
|
||
token = self._get_replicate_api_token()
|
||
if token:
|
||
env["REPLICATE_API_TOKEN"] = token
|
||
return env
|
||
|
||
|
||
def _open_blender_with_scene(self, glb_path, hdri_path):
|
||
# Kopiere ggf. vorbereitete Blender-Template-Datei, z.B. "scene_template.blend"
|
||
blender_path = "/Applications/Blender.app/Contents/MacOS/Blender"
|
||
# Template: In Blender musst du ein Skript haben, das das .glb und das hdri_path lädt
|
||
# oder als Startscript mitgibst!
|
||
cmd = [
|
||
"/Applications/Blender.app/Contents/MacOS/Blender",
|
||
"--python", "scene_setup.py",
|
||
"--", glb_path, hdri_path or ""
|
||
]
|
||
try:
|
||
subprocess.Popen(cmd)
|
||
return {"status": "opened_blender"}
|
||
except Exception as e:
|
||
self._js("window.on_error", f"Failed to open in Blender: {e}")
|
||
return {"status": "blender_error"}
|
||
|
||
|
||
def edit_external(self, filepath: str):
|
||
# filepath kommt von 2D als "output/foo.png" oder von 3D als "output/foo.png"
|
||
ext = os.path.splitext(filepath)[1].lower()
|
||
full_path = os.path.join(BASE_DIR, "web", filepath)
|
||
# Ist es ein Bild oder das PNG eines 3D-Modells?
|
||
if ext == ".png":
|
||
# Prüfe, ob ein GLB daneben liegt
|
||
fbase = os.path.splitext(full_path)[0]
|
||
glb_path = fbase + ".glb"
|
||
# Prüfe, ob auch eine HDRI da ist (Sidecar JSON auslesen)
|
||
hdri_path = None
|
||
sidecar = full_path + ".json"
|
||
if os.path.isfile(sidecar):
|
||
with open(sidecar, "r", encoding="utf-8") as f:
|
||
meta = json.load(f)
|
||
if "hdri_seamless" in meta:
|
||
hdri_path = os.path.join(os.path.dirname(full_path), meta["hdri_seamless"])
|
||
# Wenn .glb existiert → ist 3D Model: Blender öffnen!
|
||
if os.path.isfile(glb_path):
|
||
return self._open_blender_with_scene(glb_path, hdri_path)
|
||
# Sonst wie bisher (Photoshop)
|
||
try:
|
||
subprocess.Popen(["open", "-a", "Adobe Photoshop 2024", full_path])
|
||
except Exception as e:
|
||
self._js("window.on_error", f"Failed to open in Photoshop: {e}")
|
||
return {"status": "opened_external"}
|
||
|
||
|
||
# ------------- Generate 3D Model -------------
|
||
def generate_3d_model(self, filepath: str):
|
||
# filepath kommt aus JS als "output/flower-1.png"
|
||
full_path = os.path.join(BASE_DIR, "web", filepath)
|
||
if not os.path.isfile(full_path):
|
||
self._js("window.on_error", f"Image not found: {full_path}")
|
||
return {"status": "missing_file"}
|
||
|
||
threading.Thread(target=self._run_generate_3d, args=(full_path,), daemon=True).start()
|
||
return {"status": "started_3d"}
|
||
|
||
def _run_generate_3d(self, img_path: str):
|
||
"""
|
||
Execute the user‐provided script `image_to_3d.py <img_path>`, wait for completion,
|
||
then, if a .glb was produced, call JS callback window.on_3d_generated(glb_path).
|
||
Otherwise, call window.on_error(…).
|
||
"""
|
||
try:
|
||
if not self._get_replicate_api_token():
|
||
self._js(
|
||
"window.on_error",
|
||
"Replicate API token missing. Open settings and add it.",
|
||
)
|
||
return
|
||
|
||
# Construct the command: sys.executable ensures we run with same Python interpreter
|
||
script_path = os.path.join(BASE_DIR, "image_to_3d.py")
|
||
cmd = [sys.executable, script_path, img_path]
|
||
base, _ = os.path.splitext(img_path)
|
||
glb_path = base + ".glb"
|
||
previous_mtime = os.path.getmtime(glb_path) if os.path.isfile(glb_path) else None
|
||
result = subprocess.run(
|
||
cmd,
|
||
capture_output=True,
|
||
text=True,
|
||
env=self._subprocess_env(),
|
||
)
|
||
|
||
# Check for success: sidecar or actual .glb in same folder
|
||
current_mtime = os.path.getmtime(glb_path) if os.path.isfile(glb_path) else None
|
||
|
||
if current_mtime is not None and current_mtime != previous_mtime:
|
||
# Notify JavaScript
|
||
self._js("window.on_3d_generated", glb_path)
|
||
else:
|
||
# If the script failed or no .glb was produced, send an error
|
||
err_msg = (
|
||
result.stderr.strip()
|
||
or result.stdout.strip()
|
||
or "3D conversion did not produce .glb"
|
||
)
|
||
self._js("window.on_error", f"3D Model failed: {err_msg}")
|
||
except Exception as e:
|
||
self._js("window.on_error", f"3D conversion exception: {e}")
|
||
|
||
|
||
|
||
|
||
|
||
# ------------- NEW: Generate 3D Model PLUS equirectangular map -------------
|
||
|
||
def generate_3d_and_hdri(self, filepath: str):
|
||
"""
|
||
Called when user wants both 3D model and equirectangular HDRI for an image.
|
||
Kicks off both processes, waits for both, and (later) calls the next step.
|
||
"""
|
||
full_path = os.path.join(BASE_DIR, "web", filepath)
|
||
if not os.path.isfile(full_path):
|
||
self._js("window.on_error", f"Image not found: {full_path}")
|
||
return {"status": "missing_file"}
|
||
# Start thread for orchestration
|
||
threading.Thread(target=self._run_3d_and_hdri, args=(full_path,), daemon=True).start()
|
||
return {"status": "started_3d_and_hdri"}
|
||
|
||
def _run_3d_and_hdri(self, img_path: str):
|
||
# 1. Get object info (from PNG metadata or sidecar)
|
||
obj_name = self._extract_object_from_json(img_path)
|
||
print(f"[3D+HDRI] Object name from JSON: {obj_name}")
|
||
if not obj_name:
|
||
self._js("window.on_error", "Could not determine object name for HDRI prompt.")
|
||
return
|
||
|
||
# 2. Generate short background prompt using Mistral (Ollama) with streaming debug
|
||
user_prompt = f"""Suggest a short, vivid background description for a 3D scene featuring {obj_name}. \
|
||
The background should be immersive and plausible, suitable as an equirectangular HDRI for 3D rendering. \
|
||
Do NOT mention the object itself. Describe the environment in a concise way, e.g.: "lush wheat field under a blue sky", "bustling medieval town square at dawn", "mysterious alien jungle with glowing plants", etc. Respond only with the description, no extra text. Only describe the appearance of the environment, no other attributes than what is to see. \
|
||
Keep it short and precise, use 2 to 5 words max! No punctuation, just commas if needed. All lowercase. As short as possible!"""
|
||
print(f"[HDRI LLM] Sending prompt to LLM:\n{user_prompt}\n")
|
||
|
||
pano_prompt = None
|
||
for chunk in stream_ollama(user_prompt):
|
||
print(f"[HDRI LLM] {chunk!r}")
|
||
pano_prompt = chunk.strip()
|
||
if not pano_prompt:
|
||
self._js("window.on_error", "Could not generate panorama prompt.")
|
||
return
|
||
|
||
# 3. Prepend required prefix and print
|
||
pano_prompt_full = "hdri panorama view, equirectangular, 3d rendering of " + pano_prompt
|
||
prompt_with_colons = f":{pano_prompt_full}:"
|
||
print(f"[HDRI PROMPT] Final panorama prompt:\n{prompt_with_colons}\n")
|
||
|
||
# 4. Prepare output path in OUTPUT_DIR
|
||
out_base = os.path.splitext(os.path.basename(img_path))[0]
|
||
hdri_filename = f"{out_base}_hdri_seamless.png"
|
||
hdri_output_path = os.path.join(OUTPUT_DIR, hdri_filename)
|
||
print(f"[HDRI PATH] Will save to: {hdri_output_path}")
|
||
|
||
# 5. Run both tasks concurrently
|
||
import concurrent.futures
|
||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||
fut_3d = executor.submit(self._run_generate_3d_and_return, img_path)
|
||
fut_hdri = executor.submit(self._run_equirect_map, prompt_with_colons, hdri_output_path)
|
||
done, _ = concurrent.futures.wait([fut_3d, fut_hdri], return_when=concurrent.futures.ALL_COMPLETED)
|
||
|
||
# 6. Check results
|
||
err3d = fut_3d.exception() if fut_3d.done() else "3D model task did not finish"
|
||
errhdri = fut_hdri.exception() if fut_hdri.done() else "HDRI task did not finish"
|
||
|
||
if err3d or errhdri:
|
||
msg = ""
|
||
if err3d:
|
||
msg += f"3D model error: {err3d}\n"
|
||
if errhdri:
|
||
msg += f"HDRI error: {errhdri}\n"
|
||
self._js("window.on_error", msg.strip())
|
||
return
|
||
|
||
glb_path = fut_3d.result()
|
||
hdri_path = fut_hdri.result()
|
||
|
||
# 7. Write HDRI path into the original image's JSON sidecar (absolute, relative, or both)
|
||
sidecar_path = img_path + ".json"
|
||
try:
|
||
if os.path.isfile(sidecar_path):
|
||
with open(sidecar_path, "r", encoding="utf-8") as f:
|
||
meta = json.load(f)
|
||
else:
|
||
meta = {}
|
||
|
||
meta["hdri_seamless"] = os.path.relpath(hdri_path, OUTPUT_DIR) # or just hdri_filename
|
||
with open(sidecar_path, "w", encoding="utf-8") as f:
|
||
json.dump(meta, f, indent=2)
|
||
print(f"[HDRI JSON] Wrote HDRI path to {sidecar_path}: {meta['hdri_seamless']}")
|
||
except Exception as e:
|
||
print(f"[HDRI JSON] ERROR writing HDRI path: {e}")
|
||
|
||
self._on_3d_and_hdri_ready(glb_path, hdri_path)
|
||
|
||
def _extract_object_from_json(self, img_path: str) -> str:
|
||
"""
|
||
Helper: Tries to read the object/positive prompt from sidecar or PNG metadata.
|
||
"""
|
||
try:
|
||
sidecar = img_path + ".json"
|
||
if os.path.isfile(sidecar):
|
||
with open(sidecar, "r", encoding="utf-8") as f:
|
||
meta = json.load(f)
|
||
return meta.get("positive prompt", "").split(",")[0]
|
||
else:
|
||
# Try embedded PNG metadata as fallback
|
||
img = Image.open(img_path)
|
||
info = img.info
|
||
if "sd_json" in info:
|
||
d = json.loads(info["sd_json"])
|
||
return d.get("positive prompt", "").split(",")[0]
|
||
except Exception:
|
||
pass
|
||
return None
|
||
|
||
def _get_panorama_prompt(self, obj_name: str) -> str:
|
||
"""
|
||
Calls Mistral (Ollama) to generate a short background panorama prompt.
|
||
"""
|
||
user_prompt = f"""Suggest a short, vivid background description for a 3D scene featuring {obj_name}. \
|
||
The background should be immersive and plausible, suitable as an equirectangular HDRI for 3D rendering. \
|
||
Do NOT mention the object itself. Describe the environment in a concise way, e.g.: "a lush wheat field under a blue sky", "a bustling medieval town square at dawn", "a mysterious alien jungle with glowing plants", etc. Respond only with the description, no extra text."""
|
||
|
||
try:
|
||
for response in stream_ollama(user_prompt):
|
||
last = response.strip()
|
||
return last
|
||
except Exception:
|
||
return None
|
||
|
||
def _run_generate_3d_and_return(self, img_path: str):
|
||
"""
|
||
Same as _run_generate_3d, but returns the GLB path on success (or raises).
|
||
"""
|
||
if not self._get_replicate_api_token():
|
||
raise RuntimeError("Replicate API token missing. Open settings and add it.")
|
||
|
||
script_path = os.path.join(BASE_DIR, "image_to_3d.py")
|
||
cmd = [sys.executable, script_path, img_path]
|
||
base, _ = os.path.splitext(img_path)
|
||
glb_path = base + ".glb"
|
||
previous_mtime = os.path.getmtime(glb_path) if os.path.isfile(glb_path) else None
|
||
result = subprocess.run(
|
||
cmd,
|
||
capture_output=True,
|
||
text=True,
|
||
env=self._subprocess_env(),
|
||
)
|
||
current_mtime = os.path.getmtime(glb_path) if os.path.isfile(glb_path) else None
|
||
if current_mtime is not None and current_mtime != previous_mtime:
|
||
self._js("window.on_3d_generated", glb_path)
|
||
return glb_path
|
||
else:
|
||
err_msg = (
|
||
result.stderr.strip()
|
||
or result.stdout.strip()
|
||
or "3D conversion did not produce .glb"
|
||
)
|
||
raise RuntimeError(err_msg)
|
||
|
||
def _run_equirect_map(self, prompt: str, output_path: str):
|
||
"""
|
||
Calls your equi map generator script.
|
||
Logs command and full output for debugging.
|
||
"""
|
||
script_path = os.path.join(BASE_DIR, "generate_equirect.py")
|
||
cmd = [sys.executable, script_path, "--prompt", prompt, "--output", output_path]
|
||
print(f"[HDRI CMD] {' '.join(cmd)}")
|
||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||
print(f"[HDRI STDOUT]\n{result.stdout}")
|
||
print(f"[HDRI STDERR]\n{result.stderr}")
|
||
if result.returncode == 0 and os.path.isfile(output_path):
|
||
print(f"[HDRI OK] Created: {output_path}")
|
||
return output_path
|
||
else:
|
||
raise RuntimeError(result.stderr.strip() or f"HDRI generation failed for {output_path}")
|
||
|
||
def _on_3d_and_hdri_ready(self, glb_path: str, hdri_path: str):
|
||
"""
|
||
Dummy function to handle successful completion of both tasks.
|
||
For now: just print/log, later: implement what you need (e.g. combine assets, notify UI, etc).
|
||
"""
|
||
print(f"3D model ready: {glb_path}")
|
||
print(f"HDRI ready: {hdri_path}")
|
||
# Example: self._js("window.on_3d_and_hdri_ready", glb_path, hdri_path)
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
# ------------- END of Generate 3D Model & Equi methods -------------
|
||
def get_initial_images(self):
|
||
entries = []
|
||
for fname in os.listdir(OUTPUT_DIR):
|
||
if not fname.lower().endswith(".png"):
|
||
continue
|
||
|
||
fpath = os.path.join(OUTPUT_DIR, fname)
|
||
# Prüfe, ob Sidecar existiert & valides JSON ist
|
||
sidecar = fpath + ".json"
|
||
if not os.path.isfile(sidecar):
|
||
continue
|
||
try:
|
||
with open(sidecar, "r", encoding="utf-8") as f:
|
||
meta = json.load(f)
|
||
# Optional: Check ob wirklich ein Prompt drin ist
|
||
if not meta.get("positive prompt"):
|
||
continue
|
||
created = int(meta.get("created", 0))
|
||
except Exception:
|
||
continue
|
||
|
||
rel = os.path.relpath(os.path.join(OUTPUT_DIR, fname), WEB_DIR)
|
||
entries.append({
|
||
"filepath": rel,
|
||
"created": created
|
||
})
|
||
|
||
return sorted(entries, key=lambda e: e["created"], reverse=True)
|
||
|
||
def get_3d_models(self):
|
||
entries = []
|
||
for fname in os.listdir(OUTPUT_DIR):
|
||
if not fname.lower().endswith(".png"):
|
||
continue
|
||
png_path = os.path.join(OUTPUT_DIR, fname)
|
||
fbase = os.path.splitext(fname)[0]
|
||
glb_path = os.path.join(OUTPUT_DIR, fbase + ".glb")
|
||
if os.path.isfile(glb_path):
|
||
# Lade "created" aus Sidecar-JSON
|
||
created = 0
|
||
sidecar = png_path + ".json"
|
||
if os.path.isfile(sidecar):
|
||
try:
|
||
with open(sidecar, "r", encoding="utf-8") as f:
|
||
meta = json.load(f)
|
||
created = int(meta.get("created", 0))
|
||
except Exception:
|
||
pass
|
||
rel_img = os.path.relpath(png_path, WEB_DIR)
|
||
rel_glb = os.path.relpath(glb_path, WEB_DIR)
|
||
entries.append({
|
||
"img": rel_img,
|
||
"glb": rel_glb,
|
||
"created": created
|
||
})
|
||
# Sortiere nach created DESC
|
||
models = sorted(entries, key=lambda e: e["created"], reverse=True)
|
||
# Entferne "created" vor dem Rückgeben (falls Frontend das nicht braucht)
|
||
for m in models:
|
||
if "created" in m:
|
||
del m["created"]
|
||
return models
|
||
|
||
def get_image_json(self, filepath: str):
|
||
# filepath kommt aus JS als "output/flower-1.png"
|
||
full_path = os.path.join(BASE_DIR, "web", filepath)
|
||
try:
|
||
# Always prefer the sidecar if available (contains most up-to-date data)
|
||
sidecar = full_path + ".json"
|
||
if os.path.isfile(sidecar):
|
||
with open(sidecar, "r", encoding="utf-8") as f:
|
||
d = json.load(f)
|
||
if "created" in d:
|
||
del d["created"]
|
||
return json.dumps(d, indent=2)
|
||
# Fallback: try embedded metadata in PNG
|
||
img = Image.open(full_path)
|
||
info = img.info
|
||
if "sd_json" in info:
|
||
d = json.loads(info["sd_json"])
|
||
if "created" in d:
|
||
del d["created"]
|
||
return json.dumps(d, indent=2)
|
||
else:
|
||
return json.dumps({"error": "No embedded generation data found."})
|
||
except Exception as e:
|
||
return json.dumps({"error": f"Failed to read metadata: {e}"})
|
||
|
||
def _monitor_folder(self):
|
||
"""
|
||
Every 2 seconds, scan OUTPUT_DIR for PNGs. If any file is new or its mtime changed,
|
||
send window.on_image_updated(filepath, dataUri) to front end.
|
||
"""
|
||
while True:
|
||
time.sleep(2)
|
||
for fname in os.listdir(OUTPUT_DIR):
|
||
if not fname.lower().endswith(".png"):
|
||
continue
|
||
fpath = os.path.join(OUTPUT_DIR, fname)
|
||
try:
|
||
mtime = os.path.getmtime(fpath)
|
||
except FileNotFoundError:
|
||
continue
|
||
previous = self._mtimes.get(fpath)
|
||
if previous is None or mtime > previous:
|
||
# New or modified
|
||
self._mtimes[fpath] = mtime
|
||
try:
|
||
data_uri = image_to_base64(fpath)
|
||
rel = os.path.relpath(fpath, WEB_DIR) # → "output/flower-2.png"
|
||
self._js("window.on_image_updated", rel, data_uri)
|
||
except Exception:
|
||
pass
|
||
|
||
def _js(self, fn: str, *args):
|
||
import json
|
||
args_json = json.dumps(list(args))
|
||
code = f"{fn}.apply(null, {args_json})"
|
||
self.window.evaluate_js(code)
|
||
|
||
|
||
# ------------- Launch PyWebview -------------
|
||
|
||
if __name__ == "__main__":
|
||
api = Api()
|
||
window = webview.create_window(
|
||
"SD 3D Model Image Gen",
|
||
"web/index.html",
|
||
js_api=api,
|
||
width=900,
|
||
height=1000,
|
||
min_size=(650, 500),
|
||
resizable=True
|
||
)
|
||
api.set_window(window)
|
||
webview.start(debug=True)
|