import io, json from typing import List, Dict, Optional, Tuple import gradio as gr import numpy as np from PIL import Image, ImageDraw import torch from diffusers import ( StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, StableDiffusionXLInpaintPipeline, StableDiffusionXLControlNetPipeline, ControlNetModel, StableDiffusionUpscalePipeline, DPMSolverMultistepScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, HeunDiscreteScheduler, ) # ---------------- Optional deps (safe imports: ไม่มีก็ข้าม) ---------------- try: from rembg import remove as rembg_remove except Exception: rembg_remove = None _HAS_GFP = False GFPGANer = None GFP = None try: import gfpgan # type: ignore if hasattr(gfpgan, "GFPGANer"): GFPGANer = gfpgan.GFPGANer # type: ignore _HAS_GFP = True except Exception as e: print("[WARN] GFPGAN not available:", e) _HAS_REALESRGAN = False RealESRGAN = None REALSR = None try: from realesrgan import RealESRGAN # type: ignore _HAS_REALESRGAN = True except Exception as e: print("[WARN] RealESRGAN not available:", e) # ---------------- Runtime setup ---------------- device = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.float16 if device == "cuda" else torch.float32 # ---------------- Registries ---------------- MODELS: List[Tuple[str,str,str]] = [ ("stabilityai/stable-diffusion-xl-base-1.0", "SDXL Base 1.0", "เอนกประสงค์"), ("stabilityai/stable-diffusion-xl-refiner-1.0","SDXL Refiner", "เสริมรายละเอียด (pass 2)"), ("SG161222/RealVisXL_V4.0", "RealVis XL v4", "โฟโต้เรียล คน/สินค้า"), ("Lykon/dreamshaper-xl-v2", "DreamShaper XL","แฟนตาซี-เรียลลิสติก"), ("RunDiffusion/Juggernaut-XL", "Juggernaut XL", "คอนทราสต์แรง"), ("emilianJR/epiCRealismXL", "EpicRealism XL","แฟชั่น/พอร์เทรต"), ("black-forest-labs/FLUX.1-dev", "FLUX.1-dev", "แนวสมัยใหม่ (ไม่ใช่ SDXL)"), ("stabilityai/sd-turbo", "SD-Turbo", "เร็วมากสำหรับร่างไอเดีย"), ("stabilityai/stable-diffusion-2-1", "SD 2.1", "แลนด์สเคปกว้าง"), ("runwayml/stable-diffusion-v1-5", "SD 1.5", "คลาสสิก"), ("timbrooks/instruct-pix2pix", "Instruct-Pix2Pix","แก้ภาพตามคำสั่ง"), ] LORAS: List[Tuple[str,str,str]] = [ ("ByteDance/SDXL-Lightning", "SDXL-Lightning", "สปีด"), ("ostris/epicrealism-xl-lora", "EpicRealism XL", "โทนจริง"), ("alpha-diffusion/sdxl-anime-lora", "Anime-Style XL", "อนิเม"), ("alpha-diffusion/sdxl-cinematic-lora","Cinematic-Drama", "แสงหนัง"), ("alpha-diffusion/sdxl-watercolor-lora","Watercolor", "สีน้ำ"), ("alpha-diffusion/sdxl-fashion-lora", "Fashion", "แฟชั่น"), ("alpha-diffusion/sdxl-product-lora", "Product-Studio", "สินค้า"), ("alpha-diffusion/sdxl-interior-lora", "Interior-Archi", "สถาปัตย์"), ("alpha-diffusion/sdxl-food-lora", "Food-Tasty", "อาหาร"), ("alpha-diffusion/sdxl-logo-lora", "Logo-Clean", "โลโก้"), ] # ใช้ 5 ชนิดหลักเพื่อ UI กระชับและเสถียร CONTROLNETS: List[Tuple[str,str,str,str]] = [ ("diffusers/controlnet-canny-sdxl-1.0", "Canny", "เส้นขอบ", "canny"), ("diffusers/controlnet-openpose-sdxl-1.0", "OpenPose", "ท่าทางคน", "pose"), ("diffusers/controlnet-depth-sdxl-1.0", "Depth", "ระยะลึก", "depth"), ("diffusers/controlnet-softedge-sdxl-1.0", "SoftEdge", "เส้นนุ่ม", "softedge"), ("diffusers/controlnet-lineart-sdxl-1.0", "Lineart", "เส้นร่าง", "lineart"), ] PRESETS = { "Cinematic": ", cinematic lighting, 50mm, bokeh, film grain, high dynamic range", "Studio": ", studio photo, softbox lighting, sharp focus, high detail", "Anime": ", anime style, clean lineart, vibrant colors, high quality", "Product": ", product photography, seamless background, diffused light, reflections", } NEG_DEFAULT = "lowres, blurry, bad anatomy, extra fingers, watermark, jpeg artifacts, text" SCHEDULERS = { "DPM-Solver (Karras)": DPMSolverMultistepScheduler, "Euler": EulerDiscreteScheduler, "Euler a": EulerAncestralDiscreteScheduler, "Heun": HeunDiscreteScheduler, } # ---------------- Caches ---------------- PIPE_CACHE: Dict[str, object] = {} CONTROL_CACHE: Dict[str, ControlNetModel] = {} UPSCALE_PIPE: Optional[StableDiffusionUpscalePipeline] = None # ---------------- Helpers ---------------- def set_sched(pipe, name: str): cls = SCHEDULERS.get(name, DPMSolverMultistepScheduler) pipe.scheduler = cls.from_config(pipe.scheduler.config) def seed_gen(sd: int): if sd is None or sd < 0: return None g = torch.Generator(device=("cuda" if device=="cuda" else "cpu")) g.manual_seed(int(sd)); return g def prep_pipe(model_id: str, control_ids: List[str]): key = f"{model_id}|{'-'.join(control_ids) if control_ids else 'none'}" if key in PIPE_CACHE: return PIPE_CACHE[key] if control_ids: cn_models = [] for cid in control_ids: if cid not in CONTROL_CACHE: CONTROL_CACHE[cid] = ControlNetModel.from_pretrained(cid, torch_dtype=dtype, use_safetensors=True) cn_models.append(CONTROL_CACHE[cid]) pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model_id, controlnet=cn_models, torch_dtype=dtype, use_safetensors=True) else: pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=dtype, use_safetensors=True) pipe.to(device) try: if device == "cuda": pipe.enable_vae_tiling(); pipe.enable_vae_slicing() pipe.enable_xformers_memory_efficient_attention() else: pipe.enable_attention_slicing() except Exception: pass PIPE_CACHE[key] = pipe return pipe def apply_loras(pipe, lora_ids: List[str]): for rid in [x for x in lora_ids if x]: try: pipe.load_lora_weights(rid) except Exception as e: print(f"[LoRA] load failed {rid}: {e}") def to_info(meta: dict) -> str: return json.dumps(meta, ensure_ascii=False, indent=2) # ---------------- Post-process ---------------- def ensure_upscalers(): global UPSCALE_PIPE, GFP, REALSR if UPSCALE_PIPE is None: try: UPSCALE_PIPE = StableDiffusionUpscalePipeline.from_pretrained( "stabilityai/stable-diffusion-x4-upscaler", torch_dtype=dtype, use_safetensors=True ).to(device) except Exception as e: print("[Upscaler] SD x4 not available:", e) if _HAS_GFP and GFP is None and GFPGANer is not None: try: GFP = GFPGANer(model_path=None, upscale=1, arch="clean", channel_multiplier=2) except Exception as e: print("[GFPGAN] init failed:", e) if _HAS_REALESRGAN and REALSR is None and device == "cuda": try: REALSR = RealESRGAN(torch.device("cuda"), scale=4) # ต้องมี weights เองจึงจะทำงานจริง except Exception as e: REALSR = None print("[RealESRGAN] init failed:", e) def post_process(img: Image.Image, do_up: bool, do_face: bool, do_bg: bool): ensure_upscalers() out = img # Upscale: RealESRGAN (ถ้ามี) > SD x4 > skip if do_up: try: if REALSR is not None: out = Image.fromarray(REALSR.predict(np.array(out))) elif UPSCALE_PIPE is not None: if device == "cuda": with torch.autocast("cuda"): out = UPSCALE_PIPE(prompt="", image=out).images[0] else: out = UPSCALE_PIPE(prompt="", image=out).images[0] except Exception as e: print("[Upscale] skipped:", e) if do_face and _HAS_GFP and GFP is not None: try: _, _, restored = GFP.enhance(np.array(out), has_aligned=False, only_center_face=False, paste_back=True) out = Image.fromarray(restored) except Exception as e: print("[GFPGAN] skipped:", e) if do_bg and rembg_remove is not None: try: out = Image.open(io.BytesIO(rembg_remove(np.array(out)))) except Exception as e: print("[rembg] skipped:", e) return out # ---------------- Generators ---------------- def run_txt2img( model_id, model_custom, prompt, preset, negative, steps, cfg, width, height, scheduler_name, seed, lora_selected, lora_custom, ctrl_selected, img_canny, img_pose, img_depth, img_softedge, img_lineart, do_up, do_face, do_bg ): if not prompt or not str(prompt).strip(): raise gr.Error("กรุณากรอก prompt") model = (model_custom.strip() or model_id).strip() if preset and preset in PRESETS: prompt = prompt + PRESETS[preset] if not negative or not negative.strip(): negative = NEG_DEFAULT # ControlNet mapping (เฉพาะภาพที่อัปโหลดจริง) label_to_img = { "Canny": img_canny, "OpenPose": img_pose, "Depth": img_depth, "SoftEdge": img_softedge, "Lineart": img_lineart } control_ids, cond_images = [], [] for cid, label, note, key in CONTROLNETS: if label in ctrl_selected and label_to_img.get(label) is not None: control_ids.append(cid); cond_images.append(label_to_img[label]) pipe = prep_pipe(model, control_ids) set_sched(pipe, scheduler_name) # LoRA lora_ids = [s.split(" — ")[0].strip() for s in (lora_selected or [])] if lora_custom and lora_custom.strip(): lora_ids += [x.strip() for x in lora_custom.split(",") if x.strip()] apply_loras(pipe, lora_ids) width, height = int(width), int(height) gen = seed_gen(seed) if device == "cuda": with torch.autocast("cuda"): if control_ids: img = pipe( prompt=prompt, negative_prompt=negative, width=width, height=height, num_inference_steps=int(steps), guidance_scale=float(cfg), controlnet_conditioning_image=cond_images if len(cond_images)>1 else cond_images[0], generator=gen ).images[0] else: img = pipe( prompt=prompt, negative_prompt=negative, width=width, height=height, num_inference_steps=int(steps), guidance_scale=float(cfg), generator=gen ).images[0] else: if control_ids: img = pipe( prompt=prompt, negative_prompt=negative, width=width, height=height, num_inference_steps=int(steps), guidance_scale=float(cfg), controlnet_conditioning_image=cond_images if len(cond_images)>1 else cond_images[0], generator=gen ).images[0] else: img = pipe( prompt=prompt, negative_prompt=negative, width=width, height=height, num_inference_steps=int(steps), guidance_scale=float(cfg), generator=gen ).images[0] img = post_process(img, do_up, do_face, do_bg) meta = { "mode":"txt2img","model":model,"loras":lora_ids,"controlnets":ctrl_selected, "prompt":prompt,"negative":negative,"size":f"{width}x{height}", "steps":steps,"cfg":cfg,"scheduler":scheduler_name,"seed":seed, "post":{"upscale":do_up,"face_restore":do_face,"remove_bg":do_bg} } return img, to_info(meta) def run_img2img( model_id, model_custom, init_image, strength, prompt, preset, negative, steps, cfg, width, height, scheduler_name, seed, do_up, do_face, do_bg ): if init_image is None: raise gr.Error("โปรดอัปโหลดภาพเริ่มต้น") model = (model_custom.strip() or model_id).strip() if preset and preset in PRESETS: prompt = prompt + PRESETS[preset] if not negative or not negative.strip(): negative = NEG_DEFAULT pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(model, torch_dtype=dtype, use_safetensors=True).to(device) try: if device=="cuda": pipe.enable_xformers_memory_efficient_attention() except Exception: pass set_sched(pipe, scheduler_name); gen = seed_gen(seed) if device=="cuda": with torch.autocast("cuda"): img = pipe(prompt=prompt, negative_prompt=negative, image=init_image, strength=float(strength), num_inference_steps=int(steps), guidance_scale=float(cfg), generator=gen).images[0] else: img = pipe(prompt=prompt, negative_prompt=negative, image=init_image, strength=float(strength), num_inference_steps=int(steps), guidance_scale=float(cfg), generator=gen).images[0] img = post_process(img, do_up, do_face, do_bg) meta = {"mode":"img2img","model":model,"prompt":prompt,"neg":negative, "steps":steps,"cfg":cfg,"seed":seed,"strength":strength} return img, to_info(meta) def expand_canvas_for_outpaint(img: Image.Image, expand_px: int, direction: str) -> Tuple[Image.Image, Image.Image]: w, h = img.size if direction == "left": new = Image.new("RGBA",(w+expand_px,h),(0,0,0,0)); new.paste(img,(expand_px,0)) mask = Image.new("L",(w+expand_px,h),0); d=ImageDraw.Draw(mask); d.rectangle([0,0,expand_px,h], fill=255) elif direction == "right": new = Image.new("RGBA",(w+expand_px,h),(0,0,0,0)); new.paste(img,(0,0)) mask = Image.new("L",(w+expand_px,h),0); d=ImageDraw.Draw(mask); d.rectangle([w,0,w+expand_px,h], fill=255) elif direction == "top": new = Image.new("RGBA",(w,h+expand_px),(0,0,0,0)); new.paste(img,(0,expand_px)) mask = Image.new("L",(w,h+expand_px),0); d=ImageDraw.Draw(mask); d.rectangle([0,0,w,expand_px], fill=255) else: new = Image.new("RGBA",(w,h+expand_px),(0,0,0,0)); new.paste(img,(0,0)) mask = Image.new("L",(w,h+expand_px),0); d=ImageDraw.Draw(mask); d.rectangle([0,h,w,h+expand_px], fill=255) return new.convert("RGB"), mask def run_inpaint_outpaint( model_id, model_custom, base_image, mask_image, mode, expand_px, expand_dir, prompt, preset, negative, steps, cfg, width, height, scheduler_name, seed, strength, do_up, do_face, do_bg ): if base_image is None: raise gr.Error("โปรดอัปโหลดภาพฐาน") model = (model_custom.strip() or model_id).strip() if preset and preset in PRESETS: prompt = prompt + PRESETS[preset] if not negative or not negative.strip(): negative = NEG_DEFAULT pipe = StableDiffusionXLInpaintPipeline.from_pretrained(model, torch_dtype=dtype, use_safetensors=True).to(device) try: if device=="cuda": pipe.enable_xformers_memory_efficient_attention() except Exception: pass set_sched(pipe, scheduler_name); gen = seed_gen(seed) if mode == "Outpaint": base_image, mask_image = expand_canvas_for_outpaint(base_image, int(expand_px), expand_dir) if device=="cuda": with torch.autocast("cuda"): img = pipe(prompt=prompt, negative_prompt=negative, image=base_image, mask_image=mask_image, strength=float(strength), num_inference_steps=int(steps), guidance_scale=float(cfg), generator=gen).images[0] else: img = pipe(prompt=prompt, negative_prompt=negative, image=base_image, mask_image=mask_image, strength=float(strength), num_inference_steps=int(steps), guidance_scale=float(cfg), generator=gen).images[0] img = post_process(img, do_up, do_face, do_bg) meta = {"mode":mode,"model":model,"prompt":prompt,"steps":steps,"cfg":cfg,"seed":seed} return img, to_info(meta) # ---------------- UI ---------------- def build_ui(): with gr.Blocks(theme=gr.themes.Soft(), title="Masterpiece SDXL Studio Pro") as demo: gr.Markdown("# 🖼️ Masterpiece SDXL Studio Pro") gr.Markdown("Text2Img • Img2Img • Inpaint/Outpaint • Multi-LoRA • ControlNet • Upscale/FaceRestore/RemoveBG (optional)") # Common controls model_dd = gr.Dropdown(choices=[m[0] for m in MODELS], value=MODELS[0][0], label="Model") model_custom = gr.Textbox(label="Custom Model ID", placeholder="(ถ้าอยากใช้โมเดลของคุณเอง กรอกที่นี่)") preset = gr.Dropdown(choices=list(PRESETS.keys()), value=None, label="Style Preset (optional)") negative = gr.Textbox(value=NEG_DEFAULT, label="Negative Prompt") steps = gr.Slider(10, 60, 30, step=1, label="Steps") cfg = gr.Slider(1.0, 12.0, 7.0, step=0.1, label="CFG") width = gr.Slider(512, 1024, 832, step=64, label="Width") height= gr.Slider(512, 1024, 832, step=64, label="Height") scheduler = gr.Dropdown(list(SCHEDULERS.keys()), value="DPM-Solver (Karras)", label="Scheduler") seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)") # LoRA & ControlNet lora_sel = gr.CheckboxGroup(choices=[f"{rid} — {lbl} ({note})" for rid,lbl,note in LORAS], label="LoRA (เลือกได้หลายตัว)") lora_custom = gr.Textbox(label="Custom LoRA IDs (comma separated)") ctrl_sel = gr.CheckboxGroup(choices=[c[1] for c in CONTROLNETS], label="ControlNet ชนิดที่ใช้") img_canny = gr.Image(type="pil", label="Canny") img_pose = gr.Image(type="pil", label="OpenPose") img_depth = gr.Image(type="pil", label="Depth") img_softedge = gr.Image(type="pil", label="SoftEdge") img_lineart = gr.Image(type="pil", label="Lineart") with gr.Row(): do_up = gr.Checkbox(False, label="Upscale x4 (ถ้ามี)") do_face = gr.Checkbox(False, label="Face Restore (ถ้ามี)") do_bg = gr.Checkbox(False, label="Remove BG (ถ้ามี)") with gr.Tab("Text → Image"): prompt_txt = gr.Textbox(lines=3, label="Prompt") btn_txt = gr.Button("🚀 Generate") out_img_txt = gr.Image(type="pil", label="Result") out_meta_txt = gr.Textbox(label="Metadata", lines=10) with gr.Tab("Image → Image"): init_img = gr.Image(type="pil", label="Init Image") strength = gr.Slider(0.1, 1.0, 0.7, 0.05, label="Strength") prompt_i2i = gr.Textbox(lines=3, label="Prompt") btn_i2i = gr.Button("🚀 Img2Img") out_img_i2i = gr.Image(type="pil", label="Result") out_meta_i2i = gr.Textbox(label="Metadata", lines=10) with gr.Tab("Inpaint / Outpaint"): base_img = gr.Image(type="pil", label="Base Image") mask_img = gr.Image(type="pil", label="Mask (ขาว=แก้, ดำ=คงเดิม)") mode_io = gr.Radio(["Inpaint","Outpaint"], value="Inpaint", label="Mode") expand_px = gr.Slider(64, 1024, 256, 64, label="Outpaint pixels") expand_dir = gr.Radio(["left","right","top","bottom"], value="right", label="Outpaint direction") prompt_io = gr.Textbox(lines=3, label="Prompt") btn_io = gr.Button("🚀 Inpaint/Outpaint") out_img_io = gr.Image(type="pil", label="Result") out_meta_io = gr.Textbox(label="Metadata", lines=10) # Bindings btn_txt.click( fn=run_txt2img, inputs=[ model_dd, model_custom, prompt_txt, preset, negative, steps, cfg, width, height, scheduler, seed, lora_sel, lora_custom, ctrl_sel, img_canny, img_pose, img_depth, img_softedge, img_lineart, do_up, do_face, do_bg ], outputs=[out_img_txt, out_meta_txt], api_name="txt2img" ) btn_i2i.click( fn=run_img2img, inputs=[ model_dd, model_custom, init_img, strength, prompt_i2i, preset, negative, steps, cfg, width, height, scheduler, seed, do_up, do_face, do_bg ], outputs=[out_img_i2i, out_meta_i2i], api_name="img2img" ) btn_io.click( fn=run_inpaint_outpaint, inputs=[ model_dd, model_custom, base_img, mask_img, mode_io, expand_px, expand_dir, prompt_io, preset, negative, steps, cfg, width, height, scheduler, seed, strength, do_up, do_face, do_bg ], outputs=[out_img_io, out_meta_io], api_name="inpaint_outpaint" ) gr.Markdown("ℹ️ ถ้าโมดูลเสริมหรือบางโมเดลไม่พร้อมใช้งาน ระบบจะข้ามอย่างปลอดภัยและแจ้งเตือนใน Console") return demo demo = build_ui() demo.queue(max_size=8).launch()