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 Scandinavian-style lounge with minimalist design and cozy atmosphere, "
    "bright open space with white-painted walls and large windows, natural light flooding in, "
    "light oak wood flooring, functional furniture with clean lines, "
    "soft neutral tones like beige, gray, and muted pastels, "
    "a low-profile sofa with textured cushions, woven wool throws, "
    "modern coffee table in natural wood, built-in shelving with simple decor like ceramic vases, "
    "indoor plants, and framed art prints, pendant lighting with matte finishes, "
    "warm ambient light, hygge-inspired touches, uncluttered layout, "
    "ultra photorealistic, 8k, crisp focus, soft lighting, Nordic aesthetic with a modern twist"
)


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()