import os import random import gradio as gr from huggingface_hub import login, hf_hub_download import spaces import torch from diffusers import DiffusionPipeline import hashlib import pickle import yaml import logging # Set up logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') # Load config file with open('config.yaml', 'r') as file: config = yaml.safe_load(file) # Authenticate using the token stored in Hugging Face Spaces secrets if 'HF_TOKEN' in os.environ: login(token=os.environ['HF_TOKEN']) logging.info("Successfully logged in with HF_TOKEN") else: logging.warning("HF_TOKEN not found in environment variables. Some functionality may be limited.") # Correctly access the config values process_config = config['config']['process'][0] # Assuming the first process is the one we want base_model = "black-forest-labs/FLUX.1-dev" lora_model = "sagar007/sagar_flux" # This isn't in the config, so we're keeping it as is trigger_word = process_config['trigger_word'] logging.info(f"Base model: {base_model}") logging.info(f"LoRA model: {lora_model}") logging.info(f"Trigger word: {trigger_word}") # Global variables pipe = None cache = {} CACHE_FILE = "image_cache.pkl" # Example prompts example_prompts = [ "Photos of sagar as superman flying in the sky, cape billowing in the wind, sagar", "Professional photo of sagar for LinkedIn headshot, DSLR quality, neutral background, sagar", "Sagar as an astronaut exploring a distant alien planet, vibrant colors, sagar", "Sagar hiking in a lush green forest, sunlight filtering through the trees, sagar", "Sagar as a wizard casting a spell, magical energy swirling around, sagar", "Sagar scoring a goal in a dramatic soccer match, stadium lights shining, sagar", "Sagar as a Roman emperor addressing a crowd, wearing a toga and laurel wreath, sagar" ] def initialize_model(): global pipe if pipe is None: try: logging.info(f"Attempting to load model: {base_model}") pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.float16, use_safetensors=True) logging.info("Moving model to CUDA...") pipe = pipe.to("cuda") logging.info(f"Successfully loaded model: {base_model}") except Exception as e: logging.error(f"Error loading model {base_model}: {str(e)}") raise def load_cache(): global cache if os.path.exists(CACHE_FILE): with open(CACHE_FILE, 'rb') as f: cache = pickle.load(f) logging.info(f"Loaded {len(cache)} cached images") def save_cache(): with open(CACHE_FILE, 'wb') as f: pickle.dump(cache, f) logging.info(f"Saved {len(cache)} cached images") def get_cache_key(prompt, cfg_scale, steps, seed, width, height, lora_scale): return hashlib.md5(f"{prompt}{cfg_scale}{steps}{seed}{width}{height}{lora_scale}".encode()).hexdigest() @spaces.GPU(duration=80) def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale): global pipe, cache if randomize_seed: seed = random.randint(0, 2**32-1) cache_key = get_cache_key(prompt, cfg_scale, steps, seed, width, height, lora_scale) if cache_key in cache: logging.info("Using cached image") return cache[cache_key], seed try: logging.info(f"Starting run_lora with prompt: {prompt}") if pipe is None: logging.info("Initializing model...") initialize_model() logging.info(f"Using seed: {seed}") generator = torch.Generator(device="cuda").manual_seed(seed) full_prompt = f"{prompt} {trigger_word}" logging.info(f"Full prompt: {full_prompt}") logging.info("Starting image generation...") image = pipe( prompt=full_prompt, num_inference_steps=steps, guidance_scale=cfg_scale, width=width, height=height, generator=generator, ).images[0] logging.info("Image generation completed successfully") # Cache the generated image cache[cache_key] = image save_cache() return image, seed except Exception as e: logging.error(f"Error during generation: {str(e)}") import traceback logging.error(traceback.format_exc()) return None, seed def update_prompt(example): return example # Load cache at startup load_cache() # Pre-generate and cache example images def cache_example_images(): for prompt in example_prompts: run_lora(prompt, process_config['sample']['guidance_scale'], process_config['sample']['sample_steps'], process_config['sample']['walk_seed'], process_config['sample']['seed'], process_config['sample']['width'], process_config['sample']['height'], 0.75) # Gradio interface setup with gr.Blocks() as app: gr.Markdown("# Text-to-Image Generation with FLUX (ZeroGPU)") with gr.Row(): with gr.Column(): prompt = gr.Textbox(label="Prompt") example_dropdown = gr.Dropdown(choices=example_prompts, label="Example Prompts") run_button = gr.Button("Generate") with gr.Column(): result = gr.Image(label="Result") with gr.Row(): cfg_scale = gr.Slider(minimum=1, maximum=20, value=process_config['sample']['guidance_scale'], step=0.1, label="CFG Scale") steps = gr.Slider(minimum=1, maximum=100, value=process_config['sample']['sample_steps'], step=1, label="Steps") with gr.Row(): width = gr.Slider(minimum=128, maximum=1024, value=process_config['sample']['width'], step=64, label="Width") height = gr.Slider(minimum=128, maximum=1024, value=process_config['sample']['height'], step=64, label="Height") with gr.Row(): seed = gr.Number(label="Seed", value=process_config['sample']['seed'], precision=0) randomize_seed = gr.Checkbox(label="Randomize seed", value=process_config['sample']['walk_seed']) lora_scale = gr.Slider(minimum=0, maximum=1, value=0.75, step=0.01, label="LoRA Scale") example_dropdown.change(update_prompt, inputs=[example_dropdown], outputs=[prompt]) run_button.click( run_lora, inputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale], outputs=[result, seed] ) # Launch the app if __name__ == "__main__": logging.info("Starting the Gradio app...") logging.info("Pre-generating example images...") cache_example_images() app.launch(share=True) logging.info("Gradio app launched successfully")