import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
import os
from safetensors import safe_open

# Detect device
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32

print(f"Using device: {device}")

# Load base model
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=dtype
)

# Try to load LoRA weights with multiple file name options
lora_loaded = False
lora_path = "./lora"

print("🔍 Looking for LoRA files...")

if os.path.exists(lora_path):
    files = os.listdir(lora_path)
    print(f"📁 Files in lora folder: {files}")
    
    # List of possible LoRA filenames to try
    possible_names = [
        "adapter_model.safetensors",
        "pytorch_lora_weights.safetensors", 
        "lora_weights.safetensors"
    ]
    
    lora_file_found = None
    for filename in possible_names:
        if filename in files:
            lora_file_found = filename
            print(f"✅ Found LoRA file: {filename}")
            break
    
    if lora_file_found:
        model_path = os.path.join(lora_path, lora_file_found)
        file_size = os.path.getsize(model_path)
        print(f"📊 File size: {file_size:,} bytes ({file_size/1024/1024:.2f} MB)")
        
        # Test if the file is readable
        try:
            print("🔧 Testing file integrity...")
            with safe_open(model_path, framework="pt") as f:
                keys = list(f.keys())
                print(f"✅ File is valid! Found {len(keys)} tensors")
                
                # Now try to load it with diffusers
                print("🎯 Loading with diffusers...")
                
                if lora_file_found == "pytorch_lora_weights.safetensors":
                    # For pytorch_lora_weights.safetensors, try loading it directly
                    pipe.load_lora_weights(lora_path, weight_name=lora_file_found)
                else:
                    # For standard naming, load normally  
                    pipe.load_lora_weights(lora_path)
                
                lora_loaded = True
                print("✅ LoRA weights loaded successfully!")
                
        except Exception as e:
            print(f"❌ Error with {lora_file_found}: {e}")
            print("💡 File appears to be corrupted or incompatible")
    else:
        print("❌ No LoRA files found with expected names")
        
else:
    print("❌ LoRA folder not found")

# Move to device
pipe = pipe.to(device)

# Basic CPU optimizations
if device == "cpu":
    try:
        pipe.enable_attention_slicing()
        print("✅ Basic CPU optimizations enabled")
    except Exception as e:
        print(f"⚠️ Could not enable CPU optimizations: {e}")

def generate(prompt, quality):
    if quality == "Fast":
        steps = 20
        guidance_scale = 7.0
    else:  # High Quality
        steps = 40
        guidance_scale = 8.5
    
    with torch.no_grad():
        image = pipe(prompt, num_inference_steps=steps, guidance_scale=guidance_scale).images[0]
    return image

# Build Gradio UI
title = "Fine-tuned Stable Diffusion"
if lora_loaded:
    title += " with LoRA ✨"
    description = "✅ LoRA weights loaded! Your custom trained model is active."
else:
    title += " (Base Model)"
    description = "⚠️ Running with base model only. Check logs for details."

description += "\nChoose 'Fast' for quicker generation or 'High Quality' for better details."

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Enter your prompt", placeholder="A beautiful artwork..."),
        gr.Dropdown(["Fast", "High Quality"], value="Fast", label="Generation Mode")
    ],
    outputs=gr.Image(label="Generated Image"),
    title=title,
    description=description
)

if __name__ == "__main__":
    demo.launch()