app.py

import gradio as gr
import spaces
import torch
from PIL import Image
import numpy as np
import cv2
import mediapipe as mp
from diffusers import StableDiffusionInpaintPipeline
from diffusers.loaders import AttnProcsLayers

def create_mask(image_pil):
    image_cv2 = np.array(image_pil)
    original_size = image_pil.size
    width, height = original_size
    mp_pose_module = mp.solutions.pose
    pose = mp_pose_module.Pose(static_image_mode=True)
    results = pose.process(cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR))
    mask = np.zeros(image_cv2.shape[:2], dtype=np.uint8)
    if results.pose_landmarks:
        landmarks = results.pose_landmarks.landmark
        # Omuz ve kalça noktalarını al
        l_shoulder = landmarks[mp_pose_module.PoseLandmark.LEFT_SHOULDER]
        r_shoulder = landmarks[mp_pose_module.PoseLandmark.RIGHT_SHOULDER]
        l_hip = landmarks[mp_pose_module.PoseLandmark.LEFT_HIP]
        r_hip = landmarks[mp_pose_module.PoseLandmark.RIGHT_HIP]
        # Göğüs merkez noktası
        cx = int((l_shoulder.x + r_shoulder.x) / 2 * width)
        # Yüksekliği omuz ile kalça arası mesafenin üst kısmına yerleştir
        y_top = int((l_shoulder.y + r_shoulder.y) / 2 * height)
        y_bottom = int((l_hip.y + r_hip.y) / 2 * height)
        cy = int((y_top * 2 + y_bottom) / 3)  # Omuz ile kalça arasının üst kısmında, göğüs hizasında
        # Maskenin genişliğini iki omuz arası mesafe olarak al
        w = int(abs(l_shoulder.x - r_shoulder.x) * width * 1.0)
        h = int((y_bottom - y_top) * 0.4)  # Omuz-kalça arası mesafenin %40'ı kadar yüksek
        # Elips çiz (daha lokal, sadece göğüs)
        cv2.ellipse(mask, (cx, cy), (w//2, h//2), 0, 0, 360, 255, -1)
    return Image.fromarray(mask).convert("RGB")

# --- Mask oluşturma fonksiyonu ---
def create_mask(image_pil):
    image_cv2 = np.array(image_pil)
    original_size = image_pil.size
    width, height = original_size
    mp_pose_module = mp.solutions.pose
    pose = mp_pose_module.Pose(static_image_mode=True)
    results = pose.process(cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR))
    mask = np.zeros(image_cv2.shape[:2], dtype=np.uint8)
    if results.pose_landmarks:
        landmarks = results.pose_landmarks.landmark
        # Omuz ve kalça noktalarını al
        l_shoulder = landmarks[mp_pose_module.PoseLandmark.LEFT_SHOULDER]
        r_shoulder = landmarks[mp_pose_module.PoseLandmark.RIGHT_SHOULDER]
        l_hip = landmarks[mp_pose_module.PoseLandmark.LEFT_HIP]
        r_hip = landmarks[mp_pose_module.PoseLandmark.RIGHT_HIP]
        # Göğüs merkez noktası
        cx = int((l_shoulder.x + r_shoulder.x) / 2 * width)
        # Yüksekliği omuz ile kalça arası mesafenin üst kısmına yerleştir
        y_top = int((l_shoulder.y + r_shoulder.y) / 2 * height)
        y_bottom = int((l_hip.y + r_hip.y) / 2 * height)
        cy = int((y_top * 2 + y_bottom) / 3)
        # Maskenin genişliğini iki omuz arası mesafe olarak al
        w = int(abs(l_shoulder.x - r_shoulder.x) * width * 1.0)
        h = int((y_bottom - y_top) * 0.4)
        # Elips çiz (sadece göğüs)
        cv2.ellipse(mask, (cx, cy), (w//2, h//2), 0, 0, 360, 255, -1)
    return Image.fromarray(mask).convert("RGB")


# --- Model pipeline'ı ilk başta yükle, tekrar tekrar yüklenmesin ---
def get_pipe(lora_scale=1.0):
    base_model = "runwayml/stable-diffusion-inpainting"         # Kendi model yolunu buraya yaz
    lora_path = "breast.safetensors"

    pipe = StableDiffusionInpaintPipeline.from_pretrained(
        base_model,
        torch_dtype=torch.float16
    ).to("cuda")
    
    pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
    # LoRA ağırlığını yükle
    pipe.load_lora_weights(lora_path)
    # Scale'ı uygula
    pipe.fuse_lora(lora_scale=lora_scale)
    return pipe

PROMPT_SYSTEM = (
    "Keep the original clothing, color, and style unchanged. "
    "Do not change the clothing color, style, or fabric. "
    "Do not change the lighting, background, or pose. "
    "Keep the face, body, and skin texture completely natural and realistic. "
    "Do not change anything below the chest or above the shoulders. "
    "Do not alter jewelry, hair, or makeup. "
    "Only modify the breast area, maintaining a natural and photorealistic appearance."
)

@spaces.GPU
def inpaint(image, breast_size, user_prompt):
    full_prompt = (
        PROMPT_SYSTEM +
        (" Increase the breast size naturally and realistically." if breast_size > 0 else
         " Decrease the breast size naturally and realistically.")
    )
    # Kullanıcı promptu varsa ekle
    if user_prompt and user_prompt.strip():
        full_prompt += " " + user_prompt.strip()

    original_size = image.size
    mask_image = create_mask(image).convert("L")

    pipe = get_pipe(lora_scale=breast_size)
    result = pipe(
        prompt=full_prompt,
        image=image,
        mask_image=mask_image,
        num_inference_steps=50,
        guidance_scale=6.0
    ).images[0]
    edited_resized = result.resize(original_size, resample=Image.LANCZOS)

    image_np = np.array(image)
    edited_np = np.array(edited_resized)
    mask_np = np.array(mask_image)
    mask_bin = (mask_np > 127).astype(np.uint8)[..., None]
    final_np = edited_np * mask_bin + image_np * (1 - mask_bin)
    final_image = Image.fromarray(final_np.astype(np.uint8))
    return final_image

def lora_selector(label="Select change"):
    return gr.Radio(
        choices=[("Bigger", 0.3), ("Smaller", -0.3)],
        value=0.3,
        label=label,
        interactive=True
    )

gr.Interface(
    fn=inpaint,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        lora_selector(label="Breast Size Change"),
        gr.Textbox(label="prompt (optional)", placeholder="")
    ],
    outputs=gr.Image(type="pil"),
    title="MVP"
).launch()