# 1. ALWAYS IMPORT SPACES ABSOLUTELY FIRST
import spaces 
import os
import torch
from diffusers import DiffusionPipeline
import gradio as gr

# 2. Retrieve token from space secrets
HF_TOKEN = os.getenv("HF_TOKEN")

# 3. Optimize the pipeline instantiation for RAM limits
model_id = "black-forest-labs/FLUX.1-dev"
pipe = DiffusionPipeline.from_pretrained(
    model_id, 
    torch_dtype=torch.bfloat16,     # Use half-precision for memory savings
    low_cpu_mem_usage=True,         # Prevent loading full weights into CPU system RAM at once
    token=HF_TOKEN
)

# 4. Mount and fuse the LoRA weights cleanly (Requires 'peft' in requirements.txt)
pipe.load_lora_weights("strangerzonehf/Flux-Icon-Kit-LoRA")
pipe.fuse_lora()

# 5. Zero GPU dynamic invocation container wrapper
@spaces.GPU(duration=60) 
def generate_icon(prompt, num_inference_steps=28, guidance_scale=3.5):
    # Dynamically move tensors onto the dynamically provisioned Zero GPU
    pipe.to("cuda")
    
    full_prompt = f"Icon Kit, {prompt}, minimalist UI UX design element, flat line icon, uniform stroke, solid background"
    
    # Run prediction
    output = pipe(
        prompt=full_prompt,
        num_inference_steps=int(num_inference_steps),
        guidance_scale=float(guidance_scale)
    )
    
    # Extract the exact first PIL Image item from the output list safely
    image = output.images[0]
    
    # CRUCIAL: Immediately dump weights back to CPU to cleanly yield the Zero GPU slot
    pipe.to("cpu")
    return image

# 6. Expose API Endpoints using Gradio Interface
demo = gr.Interface(
    fn=generate_icon,
    inputs=[
        gr.Textbox(label="Icon Subject (e.g., 'a settings gear wheel')", placeholder="Enter icon concept..."),
        gr.Slider(minimum=15, maximum=40, value=28, step=1, label="Inference Steps"),
        gr.Slider(minimum=1.0, maximum=10.0, value=3.5, step=0.5, label="Guidance Scale")
    ],
    outputs=gr.Image(type="pil", label="Generated UI Icon Component"),
    title="Flux UI/UX Icon Kit Generator"
)

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