Update app.py

app.py

@@ -7,6 +7,11 @@ import torch
 from diffusers import DiffusionPipeline
 import hashlib
 import pickle
+import yaml
+
+# 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:
@@ -14,9 +19,9 @@ if 'HF_TOKEN' in os.environ:
 else:
     raise ValueError("HF_TOKEN not found in environment variables. Please add it to your Space's secrets.")
 
-base_model = "black-forest-labs/FLUX.1-dev"
-lora_model = "sagar007/sagar_flux"
-trigger_word = "sagar"
+base_model = config['config']['model']['name_or_path']
+lora_model = "sagar007/sagar_flux"  # This isn't in the config, so we're keeping it as is
+trigger_word = config['config']['trigger_word']
 
 # Global variables
 pipe = None
@@ -67,29 +72,29 @@ def get_cache_key(prompt, cfg_scale, steps, seed, width, height, lora_scale):
 @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:
         print("Using cached image")
         return cache[cache_key], seed
-    
+
     try:
         print(f"Starting run_lora with prompt: {prompt}")
         if pipe is None:
             print("Initializing model...")
             initialize_model()
-        
+
         print(f"Using seed: {seed}")
-        
+
         generator = torch.Generator(device="cuda").manual_seed(seed)
-        
+
         full_prompt = f"{prompt} {trigger_word}"
         print(f"Full prompt: {full_prompt}")
-        
+
         print("Starting image generation...")
         image = pipe(
             prompt=full_prompt,
@@ -100,11 +105,11 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
             generator=generator,
         ).images[0]
         print("Image generation completed successfully")
-        
+
         # Cache the generated image
         cache[cache_key] = image
         save_cache()
-        
+
         return image, seed
     except Exception as e:
         print(f"Error during generation: {str(e)}")
@@ -121,7 +126,9 @@ load_cache()
 # Pre-generate and cache example images
 def cache_example_images():
     for prompt in example_prompts:
-        run_lora(prompt, 4, 20, False, 42, 1024, 1024, 0.75)
+        run_lora(prompt, config['config']['sample']['guidance_scale'], config['config']['sample']['sample_steps'], 
+                 config['config']['sample']['walk_seed'], config['config']['sample']['seed'], 
+                 config['config']['sample']['width'], config['config']['sample']['height'], 0.75)
 
 # Gradio interface setup
 with gr.Blocks() as app:
@@ -134,18 +141,18 @@ with gr.Blocks() as app:
         with gr.Column():
             result = gr.Image(label="Result")
     with gr.Row():
-        cfg_scale = gr.Slider(minimum=1, maximum=20, value=4, step=0.1, label="CFG Scale")
-        steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Steps")
+        cfg_scale = gr.Slider(minimum=1, maximum=20, value=config['config']['sample']['guidance_scale'], step=0.1, label="CFG Scale")
+        steps = gr.Slider(minimum=1, maximum=100, value=config['config']['sample']['sample_steps'], step=1, label="Steps")
     with gr.Row():
-        width = gr.Slider(minimum=128, maximum=1024, value=1024, step=64, label="Width")
-        height = gr.Slider(minimum=128, maximum=1024, value=1024, step=64, label="Height")
+        width = gr.Slider(minimum=128, maximum=1024, value=config['config']['sample']['width'], step=64, label="Width")
+        height = gr.Slider(minimum=128, maximum=1024, value=config['config']['sample']['height'], step=64, label="Height")
     with gr.Row():
-        seed = gr.Number(label="Seed", precision=0)
-        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+        seed = gr.Number(label="Seed", value=config['config']['sample']['seed'], precision=0)
+        randomize_seed = gr.Checkbox(label="Randomize seed", value=config['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],