import spaces # Import this first to avoid CUDA initialization issues import os import gradio as gr import json import torch import random import time from PIL import Image from diffusers import DiffusionPipeline # Define the device device = "cuda" # Use the 'waffles' environment variable as the access token hf_token = os.getenv('waffles') # Ensure the token is loaded correctly if not hf_token: raise ValueError("Hugging Face API token not found. Please set the 'waffles' environment variable.") # Load LoRAs from JSON file with open('loras.json', 'r') as f: loras = json.load(f) # Initialize the base model with authentication and specify the device # Initialize the base model with authentication and specify the device pipe = DiffusionPipeline.from_pretrained( "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, token=hf_token ).to(device) MAX_SEED = 2**32 - 1 class calculateDuration: def __init__(self, activity_name=""): self.activity_name = activity_name def __enter__(self): self.start_time = time.time() return self def __exit__(self, exc_type, exc_value, traceback): self.end_time = time.time() self.elapsed_time = self.end_time - self.start_time if self.activity_name: print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds") else: print(f"Elapsed time: {self.elapsed_time:.6f} seconds") @spaces.GPU(duration=90) def generate_images(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, num_images, progress): generator = torch.Generator(device=device).manual_seed(seed) images = [] with calculateDuration("Generating images"): for _ in range(num_images): # Generate each image image = pipe( prompt=f"{prompt} {trigger_word}", num_inference_steps=steps, guidance_scale=cfg_scale, width=width, height=height, generator=generator, joint_attention_kwargs={"scale": lora_scale}, ).images[0] images.append(image) return images def run_lora(prompt, cfg_scale, steps, selected_repo, randomize_seed, seed, width, height, lora_scale, num_images, progress=gr.Progress(track_tqdm=True)): if not selected_repo: raise gr.Error("You must select a LoRA before proceeding.") selected_lora = next((lora for lora in loras if lora["repo"] == selected_repo), None) if not selected_lora: raise gr.Error("Selected LoRA not found.") lora_path = selected_lora["repo"] trigger_word = selected_lora["trigger_word"] # Load LoRA weights with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"): if "weights" in selected_lora: pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"]) else: pipe.load_lora_weights(lora_path) # Set random seed for reproducibility with calculateDuration("Randomizing seed"): if randomize_seed: seed = random.randint(0, MAX_SEED) images = generate_images(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, num_images, progress) pipe.to("cuda") pipe.unload_lora_weights() return images, seed def update_selection(evt: gr.SelectData): index = evt.index selected_lora = loras[index] return f"Selected LoRA: {selected_lora['title']}", selected_lora["repo"] run_lora.zerogpu = True css = ''' #gen_btn{height: 100%} #title{text-align: center} #title h1{font-size: 3em; display:inline-flex; align-items:center} #title img{width: 100px; margin-right: 0.5em} #gallery .grid-wrap{height: auto; width: auto;} #gallery .gallery-item{width: 50px; height: 50px; margin: 0px;} /* Make buttons 50% height and width */ #gallery img{width: 100%; height: 100%; object-fit: cover;} /* Resize images to fit buttons */ #info_blob { background-color: #f0f0f0; border: 2px solid #ccc; padding: 10px; margin: 10px 0; text-align: center; font-size: 1.2em; font-weight: bold; color: #333; border-radius: 8px; } ''' with gr.Blocks(theme=gr.themes.Soft(), css=css) as app: title = gr.HTML( """