Update app.py

app.py

@@ -5,6 +5,8 @@ from huggingface_hub import login, hf_hub_download
 import spaces
 import torch
 from diffusers import DiffusionPipeline
+import hashlib
+import pickle
 
 # Authenticate using the token stored in Hugging Face Spaces secrets
 if 'HF_TOKEN' in os.environ:
@@ -16,8 +18,10 @@ base_model = "black-forest-labs/FLUX.1-dev"
 lora_model = "sagar007/sagar_flux"
 trigger_word = "sagar"
 
-# Global variable for the pipeline
+# Global variables
 pipe = None
+cache = {}
+CACHE_FILE = "image_cache.pkl"
 
 # Example prompts
 example_prompts = [
@@ -39,47 +43,50 @@ def initialize_model():
             print("Moving model to CUDA...")
             pipe = pipe.to("cuda")
             print(f"Successfully loaded base model: {base_model}")
-            
-            # Commenting out LoRA loading for now
-            # print("Downloading LoRA weights...")
-            # lora_path = download_lora_weights(lora_model)
-            # if lora_path:
-            #     print("Loading LoRA weights...")
-            #     pipe.load_lora_weights(lora_path)
-            #     print("Successfully loaded LoRA weights")
-            # else:
-            #     print("Failed to download LoRA weights. Continuing without LoRA.")
         except Exception as e:
             print(f"Error initializing model: {str(e)}")
             import traceback
             print(traceback.format_exc())
             raise
 
-def download_lora_weights(repo_id, filename="lora.safetensors"):
-    try:
-        lora_path = hf_hub_download(repo_id=repo_id, filename=filename)
-        print(f"Successfully downloaded LoRA weights from {repo_id}")
-        return lora_path
-    except Exception as e:
-        print(f"Error downloading LoRA weights: {str(e)}")
-        return None
+def load_cache():
+    global cache
+    if os.path.exists(CACHE_FILE):
+        with open(CACHE_FILE, 'rb') as f:
+            cache = pickle.load(f)
+    print(f"Loaded {len(cache)} cached images")
+
+def save_cache():
+    with open(CACHE_FILE, 'wb') as f:
+        pickle.dump(cache, f)
+    print(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
+    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()
         
-        if randomize_seed:
-            seed = random.randint(0, 2**32-1)
         print(f"Using seed: {seed}")
         
         generator = torch.Generator(device="cuda").manual_seed(seed)
         
-        # Include the trigger word in the prompt
         full_prompt = f"{prompt} {trigger_word}"
         print(f"Full prompt: {full_prompt}")
         
@@ -93,6 +100,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)}")
@@ -103,6 +115,14 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
 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, 4, 20, False, 42, 1024, 1024, 0.75)
+
 # Gradio interface setup
 with gr.Blocks() as app:
     gr.Markdown("# Text-to-Image Generation with FLUX (ZeroGPU)")
@@ -135,5 +155,7 @@ with gr.Blocks() as app:
 # Launch the app
 if __name__ == "__main__":
     print("Starting the Gradio app...")
+    print("Pre-generating example images...")
+    cache_example_images()
     app.launch(share=True)
     print("Gradio app launched successfully")