changes made to all the files

custom_stable_diffusion.py

@@ -340,18 +340,35 @@ def generate_with_multiple_concepts(models, config, image_processor, prompt, con
     """
     os.makedirs(output_dir, exist_ok=True)
     
-    if not concepts:
-        # Handle the case with no concept
-        # ... your existing code ...
-        # Make sure to return the PIL Image object
-        return pil_image
-    
+    # If no concepts provided, generate a standard image
+    if not concepts or len(concepts) == 0:
+        print("No concepts provided, generating standard image")
+        # Create a standard image without concepts
+        # You'll need to implement this part based on your existing code
+        # For now, return None
+        return None
+    
+    # Process each concept
     for concept in concepts:
-        # ... your existing code ...
-        # Make sure to return the PIL Image object
+        if concept is None:
+            continue
+            
+        print(f"Generating image for concept: {concept}")
+        concepts_dir = os.path.join(output_dir, concept)
+        os.makedirs(concepts_dir, exist_ok=True)
+        
+        output_path = os.path.join(concepts_dir, f"{concept}.png")
+        
+        text_processor = TextEmbeddingProcessor(models, config, image_processor, prompt)
+        
+        # Generate the image with the concept
+        pil_image = text_processor.prepare_embeddings_with_concepts(prompt, concept_name=concept, output_path=output_path)
+        print(f"Saved image to {output_path}")
+        
+        # Return the generated image
         return pil_image
     
-    # If we get here, return None
+    # If we get here (no valid concepts processed), return None
     return None
 
 def channel_loss(images, channel_idx=2, target_value=0.9):
@@ -585,25 +602,28 @@ def generate_with_concept_and_color(
     
     return pil_image
 
-# Function to generate multiple concepts with color guidance
-def generate_with_multiple_concepts_and_color(
-    models, 
-    config, 
-    image_processor, 
-    prompt, 
-    concepts, 
-    output_dir="concept_images",
-    blue_loss_scale=0,
-    yellow_loss_scale=0
-):
+def generate_with_multiple_concepts_and_color(models, config, image_processor, prompt, concepts, output_dir="concept_images", blue_loss_scale=0, yellow_loss_scale=0):
     """
     Generate images using multiple concepts and color guidance
     """
     os.makedirs(output_dir, exist_ok=True)
     
+    # If no concepts provided, generate a standard image with color guidance
+    if not concepts or len(concepts) == 0:
+        print("No concepts provided, generating standard image with color guidance")
+        # Create a standard image with color guidance but without concepts
+        # You'll need to implement this part based on your existing code
+        # For now, return None
+        return None
+    
+    # Process each concept
     for concept in concepts:
-        print(f"Generating image for concept: {concept} and color guidance")
-        generate_with_concept_and_color(
+        if concept is None:
+            continue
+            
+        print(f"Generating image for concept: {concept} with color guidance")
+        # Generate the image with the concept and color guidance
+        pil_image = generate_with_concept_and_color(
             models=models,
             config=config,
             image_processor=image_processor,
@@ -613,7 +633,12 @@ def generate_with_multiple_concepts_and_color(
             blue_loss_scale=blue_loss_scale,
             yellow_loss_scale=yellow_loss_scale
         )
-
+        
+        # Return the generated image
+        return pil_image
+    
+    # If we get here (no valid concepts processed), return None
+    return None
 
 # Example usage
 if __name__ == "__main__":