import gradio as gr
import torch
from diffusers import StableDiffusionXLImg2ImgPipeline
from safetensors.torch import load_file
from PIL import Image

# --- Base SDXL model ---
BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
LORA_PATH = "studioghibli_flux_r32-v2.safetensors"

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

print("🔹 Loading SDXL base model...")
pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    BASE_MODEL,
    torch_dtype=dtype,
    use_safetensors=True,
    variant="fp16" if torch.cuda.is_available() else None,
).to(device)

# --- Apply LoRA weights ---
print("🎨 Applying Ghibli-style LoRA...")
try:
    lora_weights = load_file(LORA_PATH)
    pipe.unet.load_state_dict(lora_weights, strict=False)
    print("✅ LoRA loaded successfully.")
except Exception as e:
    print(f"⚠️ Failed to load LoRA: {e}")

# --- Ghibli-style conversion ---
def ghibli_style(image, steps=30, guidance=7.5, strength=0.6, seed=42):
    if image is None:
        raise gr.Error("Please upload an image to convert.")
    
    generator = torch.Generator(device=device).manual_seed(int(seed))
    init_image = Image.open(image).convert("RGB").resize((1024, 1024))

    prompt = "Ghibli-style art, soft lighting, painterly textures, cinematic color palette"

    result = pipe(
        prompt=prompt,
        image=init_image,
        strength=float(strength),
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        generator=generator,
    ).images[0]

    return result

# --- Gradio Interface ---
demo = gr.Interface(
    fn=ghibli_style,
    inputs=[
        gr.Image(label="Upload Image", type="filepath"),
        gr.Slider(10, 50, 30, step=1, label="Inference Steps"),
        gr.Slider(1, 15, 7.5, step=0.5, label="Guidance Scale"),
        gr.Slider(0.1, 1.0, 0.6, step=0.1, label="Style Strength"),
        gr.Number(label="Seed", value=42),
    ],
    outputs=gr.Image(label="Ghibli Style Output"),
    title="Ghibli Style Image Converter",
    description="Upload any image and transform it into a Studio Ghibli-style artwork using the Flux LoRA and SDXL model.",
)

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