Spaces:
Running on Zero
Running on Zero
File size: 2,067 Bytes
f0e2bfd e813640 fb13888 9501664 fb13888 7b53522 9092351 f0e2bfd 9092351 9501664 7b53522 f0e2bfd 9092351 fb13888 f0e2bfd 9501664 f0e2bfd 9501664 f0e2bfd 7b53522 f0e2bfd 9501664 f0e2bfd 7b53522 f0e2bfd 9501664 fb13888 7b53522 fb13888 9501664 9092351 7b53522 9092351 9501664 9092351 fb13888 9501664 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | # 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()
|