import os
import gradio as gr
import torch
import logging
import sys
from PIL import Image
import numpy as np
from diffusers import StableDiffusionImg2ImgPipeline
from peft import PeftModel

# Loglama ayarları
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger("app")

print("===== Application Startup at", torch.cuda.is_available(), "=====")

# Model ve işlem fonksiyonları
def load_model(lora_model_id="meryemarpaci/sd2base-inpainting-lora"):
    """LoRA modelini yükle"""
    print("Model yükleniyor...")
    try:
        # Baz model yükleme
        print("Baz model yükleniyor...")
        model_id = "stabilityai/stable-diffusion-2-base"
        
        pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
        )
        
        # LoRA ağırlıklarını yükleme
        print(f"LoRA ağırlıkları yükleniyor: {lora_model_id}")
        pipe.unet = PeftModel.from_pretrained(pipe.unet, lora_model_id)
        
        # GPU'ya taşıma
        if torch.cuda.is_available():
            pipe = pipe.to("cuda")
            print("Model GPU'ya yüklendi")
        else:
            print("xformers mevcut değil, varsayılan dikkat kullanılıyor:", torch.cuda.is_available())
        
        print("Model başarıyla yüklendi!")
        return pipe
    except Exception as e:
        print(f"Model yüklenirken hata: {str(e)}")
        raise e

def process_inpainting(image, mask, prompt, negative_prompt, guidance_scale, steps):
    """İnpainting işlemi gerçekleştir"""
    try:
        # Model yükleme (her seferinde yüklemek performans açısından ideal değil ama daha güvenli)
        pipe = load_model()
        
        # Görüntü ve maske kontrolü
        if image is None:
            return None, "Lütfen bir görüntü yükleyin."
        if mask is None:
            return None, "Lütfen bir maske yükleyin."
            
        # Görüntüyü ve maskeyi PIL formatına dönüştür
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        if isinstance(mask, np.ndarray):
            mask = Image.fromarray(mask)
            
        # Maskeyi siyah-beyaz yap
        mask = mask.convert("L")
            
        # İnpainting işlemi
        result = pipe(
            prompt=prompt,
            image=image,
            strength=0.8,  # İnpainting için uygun değer
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            negative_prompt=negative_prompt if negative_prompt else None,
        ).images[0]
        
        return result, "İnpainting işlemi başarıyla tamamlandı!"
    except Exception as e:
        error_msg = f"İnpainting işleminde hata: {str(e)}"
        print(error_msg)
        return None, error_msg

# Basit arayüz
title = "İnpainting LoRA Demo"
description = """
# İnpainting LoRA Modeli Demo

Bu demo, görüntülerdeki maskelenmiş bölgeleri akıllıca tamamlayan bir LoRA modelini kullanır.

**Kullanım:**
1. Bir görüntü yükleyin
2. Bir maske yükleyin (beyaz bölgeler doldurulacak)
3. Prompt girin ve "İnpainting Yap" butonuna tıklayın
"""

# Gradio 3.x ile uyumlu basit arayüz
demo = gr.Interface(
    fn=process_inpainting,
    inputs=[
        gr.Image(label="Orijinal Görüntü", type="pil"),
        gr.Image(label="Maske (Beyaz alanlar doldurulacak)", type="pil"),
        gr.Textbox(label="Prompt", value="Yüksek kaliteli bir fotoğraf"),
        gr.Textbox(label="Negatif Prompt (İsteğe Bağlı)"),
        gr.Slider(label="Guidance Scale", minimum=1.0, maximum=15.0, value=7.5, step=0.5),
        gr.Slider(label="Adım Sayısı", minimum=20, maximum=100, value=50, step=1)
    ],
    outputs=[
        gr.Image(label="Sonuç"),
        gr.Textbox(label="Durum")
    ],
    title=title,
    description=description,
    allow_flagging="never"
)

# Ana uygulama
if __name__ == "__main__":
    demo.launch()