fresh pipeline for each concept

app.py

@@ -5,6 +5,7 @@ from diffusers import StableDiffusionPipeline
 import random
 from huggingface_hub import hf_hub_download
 import os
+from transformers import CLIPTokenizer
 
 # Initialize the model
 model_id = "CompVis/stable-diffusion-v1-4"
@@ -52,18 +53,14 @@ def generate_images(prompt):
     images = []
     failed_concepts = []
     
-    # Load base pipeline
-    pipe = StableDiffusionPipeline.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16 if device == "cuda" else torch.float32
-    ).to(device)
-    
-    # Store original tokenizer and text encoder state
-    original_vocab_size = len(pipe.tokenizer)
-    original_embedding_weights = pipe.text_encoder.get_input_embeddings().weight.data.clone()
-    
     for concept in concepts:
         try:
+            # Create a fresh pipeline for each concept
+            pipe = StableDiffusionPipeline.from_pretrained(
+                model_id,
+                torch_dtype=torch.float16 if device == "cuda" else torch.float32
+            ).to(device)
+            
             # Download and load concept embedding
             embed_path = download_concept_embedding(concept)
             if embed_path is None:
@@ -94,10 +91,9 @@ def generate_images(prompt):
             
             images.append(image)
             
-            # Reset tokenizer and text encoder to original state
-            pipe.tokenizer = pipe.tokenizer.__class__.from_pretrained(model_id)
-            pipe.text_encoder.resize_token_embeddings(original_vocab_size)
-            pipe.text_encoder.get_input_embeddings().weight.data.copy_(original_embedding_weights)
+            # Clean up to free memory
+            del pipe
+            torch.cuda.empty_cache()
             
         except Exception as e:
             print(f"Error processing concept {concept}: {str(e)}")