import gradio as gr
import torch
import os
from diffusers import StableDiffusionXLImg2ImgPipeline, StableDiffusionXLPipeline

# 1. SETUP MODEL & DEVICE
model_id = "stabilityai/stable-diffusion-xl-base-1.0"
lora_id = "EmoonX/Surveillance-Ghost-Eye-NVG"

# Deteksi otomatis (penting untuk Free Tier)
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if torch.cuda.is_available() else torch.float32

# Load pipeline Text-to-Image
pipe = StableDiffusionXLPipeline.from_pretrained(
    model_id, 
    torch_dtype=dtype, 
    variant="fp16" if device == "cuda" else None, 
    use_safetensors=True
)

# Load pipeline Image-to-Image (Berbagi komponen untuk hemat RAM)
pipe_img2img = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    model_id,
    torch_dtype=dtype,
    text_encoder=pipe.text_encoder,
    tokenizer=pipe.tokenizer,
    vae=pipe.vae,
    variant="fp16" if device == "cuda" else None,
    use_safetensors=True
)

# Pengaturan memori untuk Free Tier
if device == "cuda":
    pipe.to("cuda")
    pipe_img2img.to("cuda")
else:
    pipe.enable_model_cpu_offload()
    pipe_img2img.enable_model_cpu_offload()

# 2. LOAD LORA DENGAN PENANGANAN ERROR
try:
    # Ganti 'weight_name' dengan nama file asli di repo Hugging Face Anda
    # Jika nama file di repo adalah 'pytorch_lora_weights.safetensors', biarkan kosong.
    # Jika namanya spesifik, masukkan di sini.
    pipe.load_lora_weights(lora_id, weight_name="GhostEye_NVG_Surveillance_SDXL_v1.safetensors")
    pipe_img2img.load_lora_weights(lora_id, weight_name="GhostEye_NVG_Surveillance_SDXL_v1.safetensors")
    print("LoRA Loaded Successfully")
except Exception as e:
    print(f"Gagal memuat LoRA: {e}")
    # Fallback: mencoba tanpa weight_name jika gagal
    pipe.load_lora_weights(lora_id)
    pipe_img2img.load_lora_weights(lora_id)

def generate_txt2img(prompt, lora_scale, steps):
    full_prompt = f"nvg style, surveillance-core, {prompt}"
    image = pipe(
        prompt=full_prompt, 
        cross_attention_kwargs={"scale": lora_scale},
        num_inference_steps=int(steps)
    ).images[0]
    return image

def generate_img2img(input_image, prompt, strength, lora_scale, steps):
    if input_image is None: return None
    full_prompt = f"nvg style, surveillance-core, {prompt}"
    input_image = input_image.convert("RGB").resize((768, 768))
    image = pipe_img2img(
        prompt=full_prompt,
        image=input_image,
        strength=strength,
        cross_attention_kwargs={"scale": lora_scale},
        num_inference_steps=int(steps)
    ).images[0]
    return image

# 3. INTERFACE GRADIO
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 👁️ GHOST-EYE NVG Demo")
    with gr.Tabs():
        with gr.TabItem("Text to Image"):
            with gr.Row():
                with gr.Column():
                    t2i_prompt = gr.Textbox(label="Prompt", placeholder="Shadowy figure...")
                    t2i_lora = gr.Slider(0.5, 1.5, value=1.0, label="LoRA Intensity")
                    t2i_steps = gr.Slider(10, 40, value=25, step=1, label="Steps")
                    t2i_btn = gr.Button("Generate")
                with gr.Column():
                    t2i_output = gr.Image(label="Output")
        with gr.TabItem("Image to Image"):
            with gr.Row():
                with gr.Column():
                    i2i_input = gr.Image(type="pil", label="Upload Photo")
                    i2i_prompt = gr.Textbox(label="Prompt", value="night vision, grainy")
                    i2i_strength = gr.Slider(0.1, 0.9, value=0.5, label="Strength")
                    i2i_lora = gr.Slider(0.5, 1.5, value=1.0, label="LoRA Intensity")
                    i2i_btn = gr.Button("Process")
                with gr.Column():
                    i2i_output = gr.Image(label="Result")

    t2i_btn.click(generate_txt2img, inputs=[t2i_prompt, t2i_lora, t2i_steps], outputs=t2i_output)
    i2i_btn.click(generate_img2img, inputs=[i2i_input, i2i_prompt, i2i_strength, i2i_lora, i2i_steps], outputs=i2i_output)

demo.launch()