import gradio as gr
import torch
from PIL import Image
from diffusers import AutoPipelineForImage2Image

# Load SDXL base pipeline
pipe = AutoPipelineForImage2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True
).to("cuda")

# Load your LoRA weights from Hugging Face
pipe.load_lora_weights("theoracle/sdxl-tok-lounge-lora")

# Inference function
def generate(image, prompt, negative_prompt, strength):
    image = image.resize((1024, 1024))
    result = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image=image,
        strength=strength,
        guidance_scale=15,
        num_inference_steps=30
    ).images[0]
    return result

# Default prompt
default_prompt = (
    "a blue maritime-themed TOK lounge, inspired by classic coastal elegance and deep ocean hues, "
    "navy and seafoam color palette, whitewashed wooden walls, large panoramic windows with ocean views, "
    "nautical decor like ropes, brass instruments, and ship wheels subtly integrated, "
    "linen sofas in soft ivory and deep blue, striped cushions and driftwood coffee table, "
    "open shelves with coral sculptures, vintage maritime maps and glass fishing floats, "
    "light oak flooring and woven jute rug, sheer curtains billowing in a coastal breeze, "
    "ambient natural light, serene and breezy atmosphere, luxury seaside escape with timeless nautical charm, "
    "8k ultra photorealistic, crisp focus, interior photography with coastal lighting and peaceful mood"
)

default_negative = (
    "empty rooms, clutter, mess, disorganized, visible wires, tangled cables, chaotic layout, "
    "outdated design, old furniture, vintage upholstery, faded fabric, chipped wood, "
    "dirty carpet, old carpet, boring colors, muted tones, brown overload, yellow tint, "
    "stains, dust, peeling wallpaper, worn out textures, bad lighting, overexposed, "
    "low detail, blurry, deformed, flat perspective, bad composition, watermark, text, "
    "ugly, dated, dark corners, awkward layout, excessive decoration, chaotic patterns, no new doors, no new windows, no new opening"
)

# Gradio UI
gr.Interface(
    fn=generate,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Textbox(label="Prompt", value=default_prompt, lines=6),
        gr.Textbox(label="Negative Prompt", value=default_negative, lines=4),
        gr.Slider(0.1, 1.0, value=0.5, label="Strength")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="TOK Lounge Generator (LoRA SDXL)",
    description="Upload an image and generate a stylized TOK lounge interior using a LoRA fine-tuned Stable Diffusion XL model."
).launch()