import gradio as gr
import torch
import os
from diffusers import AutoPipelineForText2Image
from huggingface_hub import snapshot_download
import logging

# Logging ayarları
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ModelHandler:
    def __init__(self):
        self.pipeline = None
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
        
    def load_model(self, progress=gr.Progress()):
        try:
            if self.pipeline is not None:
                return "Model zaten yüklü."
            
            progress(0, desc="Base model indiriliyor...")
            # Base modeli indir
            base_model_path = snapshot_download(
                repo_id="black-forest-labs/FLUX.1-dev",
                local_dir="./models/base_model",
                ignore_patterns=["*.bin", "*.onnx"] if os.path.exists("./models/base_model") else None,
                token=os.getenv("HF_TOKEN")
            )
            
            progress(0.5, desc="LoRA modeli indiriliyor...")
            # LoRA modelini indir
            lora_model_path = snapshot_download(
                repo_id="codermert/ezelll_flux",
                local_dir="./models/lora_model",
                ignore_patterns=["*.bin", "*.onnx"] if os.path.exists("./models/lora_model") else None,
                token=os.getenv("HF_TOKEN")
            )
            
            progress(0.7, desc="Pipeline oluşturuluyor...")
            # Pipeline'ı oluştur
            self.pipeline = AutoPipelineForText2Image.from_pretrained(
                base_model_path,
                torch_dtype=self.dtype,
                use_safetensors=True
            ).to(self.device)
            
            progress(0.9, desc="LoRA yükleniyor...")
            # LoRA'yı yükle
            lora_path = os.path.join(lora_model_path, "lora.safetensors")
            if os.path.exists(lora_path):
                self.pipeline.load_lora_weights(lora_path)
            else:
                return "LoRA dosyası bulunamadı!"
            
            progress(1.0, desc="Tamamlandı!")
            return "Model başarıyla yüklendi! Artık görüntü oluşturmaya hazırsınız."
        except Exception as e:
            logger.error(f"Model yükleme hatası: {str(e)}")
            return f"Model yüklenirken hata oluştu: {str(e)}"
    
    def generate_image(self, prompt, use_tok=True, progress=gr.Progress()):
        try:
            if self.pipeline is None:
                return None, "Lütfen önce modeli yükleyin!"
            
            # Eğer use_tok seçeneği işaretlendiyse, prompt'a TOK ekle
            if use_tok and "TOK" not in prompt:
                prompt = f"TOK {prompt}"
            
            progress(0.2, desc="Görüntü oluşturuluyor...")
            # Görüntü oluştur
            image = self.pipeline(
                prompt,
                num_inference_steps=30,
                guidance_scale=7.5
            ).images[0]
            
            progress(1.0, desc="Tamamlandı!")
            return image, f"Oluşturulan prompt: {prompt}"
        except Exception as e:
            logger.error(f"Görüntü oluşturma hatası: {str(e)}")
            return None, f"Hata oluştu: {str(e)}"

# Model işleyiciyi oluştur
model_handler = ModelHandler()

# Gradio arayüzü
with gr.Blocks(title="Malika - FLUX Text-to-Image") as demo:
    gr.Markdown("# Malika - FLUX.1 Text-to-Image Modeliyle Görüntü Oluşturucu")
    gr.Markdown("Bu uygulama, codermert/malikafinal modelini kullanarak metinden görüntü oluşturur.")
    
    with gr.Row():
        load_model_btn = gr.Button("Modeli Yükle", variant="primary")
        model_status = gr.Textbox(label="Model Durumu", value="Model henüz yüklenmedi")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Prompt", 
                placeholder="Görüntü için prompt yazın...",
                lines=3
            )
            tok_checkbox = gr.Checkbox(
                label="Otomatik TOK Ekle", 
                value=True, 
                info="İşaretliyse prompt'a otomatik olarak TOK ekler"
            )
            generate_btn = gr.Button("Görüntü Oluştur", variant="primary")
        
        with gr.Column():
            image_output = gr.Image(label="Oluşturulan Görüntü")
            prompt_used = gr.Textbox(label="Kullanılan Prompt")
    
    load_model_btn.click(
        fn=model_handler.load_model,
        outputs=model_status
    )
    
    generate_btn.click(
        fn=model_handler.generate_image,
        inputs=[prompt_input, tok_checkbox],
        outputs=[image_output, prompt_used]
    )
    
    gr.Markdown("""
    ## Kullanım Tavsiyeleri
    - İlk olarak "Modeli Yükle" düğmesine tıklayın (bu işlem biraz zaman alabilir)
    - Eğer model için özel bir trigger sözcüğü gerekliyse 'TOK' seçeneğini işaretli bırakın
    - Daha gerçekçi sonuçlar için "photorealistic, 8K, detailed" gibi ifadeler ekleyebilirsiniz
    - Örnek: "portrait of a woman with blue eyes, photorealistic, 8K"
    
    ## Model Bilgisi
    Bu uygulama codermert/malikafinal modelini kullanmaktadır.
    Base model: black-forest-labs/FLUX.1-dev
    """)

# Arayüzü başlat
if __name__ == "__main__":
    demo.launch(share=True)