changes made to all the files

app.py

@@ -5,6 +5,9 @@ import time
 import os
 from PIL import Image
 from custom_stable_diffusion import StableDiffusionConfig, StableDiffusionModels,ImageProcessor, generate_with_multiple_concepts,generate_with_multiple_concepts_and_color
+import sys
+import transformers
+import diffusers
 st.set_page_config(
     page_title="Butterfly Color Diffusion",
     page_icon="🦋",
@@ -17,9 +20,19 @@ if 'models' not in st.session_state:
     st.session_state.models = None
     st.session_state.config = None
 
+# Add this near the top of your app.py
+debug_mode = st.sidebar.checkbox("Debug Mode", value=True)
+
 # Function to load models
 @st.cache_resource
 def load_models():
+    if debug_mode:
+        st.write("Debug: Starting model loading")
+        st.write(f"Debug: Python version: {sys.version}")
+        st.write(f"Debug: Torch version: {torch.__version__}")
+        st.write(f"Debug: Transformers version: {transformers.__version__}")
+        st.write(f"Debug: Diffusers version: {diffusers.__version__}")
+    
     config = StableDiffusionConfig(
         height=512,
         width=512,
@@ -35,6 +48,10 @@ def load_models():
     with st.spinner("Loading Stable Diffusion models... This may take a minute."):
         models.load_models()
         models.set_timesteps()
+    
+    if debug_mode:
+        st.write(f"Debug: Models loaded successfully. Device: {config.device}")
+    
     return models, config, image_processor
 
 # Title and description
@@ -160,7 +177,7 @@ if standard_button:
             caption = f"Standard Stable Diffusion"
             if concept_name:
                 caption += f" with {concept_name} concept"
-            st.image(image, caption=caption, use_container_width=True)
+            st.image(image, caption=caption, use_column_width=True)
             st.write(f"Generation time: {end_time - start_time:.2f} seconds")
 
 # Generate color-guided image
@@ -210,7 +227,7 @@ if color_button:
             caption = f"Color-Guided Stable Diffusion"
             if concept_name:
                 caption += f" with {concept_name} concept"
-            st.image(image, caption=caption, use_container_width=True)
+            st.image(image, caption=caption, use_column_width=True)
             st.write(f"Generation time: {end_time - start_time:.2f} seconds")
 
 # Explanation section

custom_stable_diffusion.py

@@ -60,27 +60,65 @@ class StableDiffusionModels:
         self.scheduler= None
 
     def load_models(self, model_version:str="CompVis/stable-diffusion-v1-4"):
-      """
-      Load all the required models for stable diffusion.
-      """
-      # Load VAE
-      self.vae = AutoencoderKL.from_pretrained(model_version, subfolder="vae")
-      
-      # Load tokenizer and text encoder - IMPORTANT: Use the correct model
-      self.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
-      self.text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+        """
+        Load all the required models for stable diffusion.
+        """
+        try:
+            # Add cache directory to ensure files are saved in a writable location
+            cache_dir = "./model_cache"
+            os.makedirs(cache_dir, exist_ok=True)
+            
+            # Load VAE
+            self.vae = AutoencoderKL.from_pretrained(
+                model_version, 
+                subfolder="vae",
+                cache_dir=cache_dir,
+                local_files_only=False
+            )
+            
+            # Load tokenizer and text encoder with explicit cache directory
+            self.tokenizer = CLIPTokenizer.from_pretrained(
+                "openai/clip-vit-large-patch14",
+                cache_dir=cache_dir,
+                local_files_only=False
+            )
+            
+            self.text_encoder = CLIPTextModel.from_pretrained(
+                "openai/clip-vit-large-patch14",
+                cache_dir=cache_dir,
+                local_files_only=False
+            )
 
-      # Load UNet
-      self.unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")
-      
-      # Load scheduler
-      self.scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
-      
-      self.vae = self.vae.to(self.config.device)
-      self.text_encoder = self.text_encoder.to(self.config.device)
-      self.unet = self.unet.to(self.config.device)
-      print(self.config.device)
-      return self
+            # Load UNet
+            self.unet = UNet2DConditionModel.from_pretrained(
+                model_version, 
+                subfolder="unet",
+                cache_dir=cache_dir,
+                local_files_only=False
+            )
+            
+            # Load scheduler
+            self.scheduler = LMSDiscreteScheduler(
+                beta_start=0.00085, 
+                beta_end=0.012, 
+                beta_schedule="scaled_linear", 
+                num_train_timesteps=1000
+            )
+            
+            # Move models to device
+            self.vae = self.vae.to(self.config.device)
+            self.text_encoder = self.text_encoder.to(self.config.device)
+            self.unet = self.unet.to(self.config.device)
+            
+            print(f"Using device: {self.config.device}")
+            return self
+            
+        except Exception as e:
+            print(f"Error loading models: {str(e)}")
+            # Add more detailed error information
+            import traceback
+            traceback.print_exc()
+            raise
     
     def set_timesteps(self, num_inference_steps:int=None):
         """
@@ -296,24 +334,25 @@ class TextEmbeddingProcessor:
       else:
           print(f"Failed to load concept: {concept_name}")
 
-def generate_with_multiple_concepts(models, config, image_processor, prompt,concepts, output_dir="generated_images"):
-        """
-        Generate images using multiple concepts and save them in separate folders
-
-        """
-
-        os.makedirs(output_dir, exist_ok=True)
-
-        for concept in concepts:
-            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)
-
-            text_processor.prepare_embeddings_with_concepts(prompt, concept_name= concept, output_path=output_path)
-            print(f"Saved iamge to {output_path}")
+def generate_with_multiple_concepts(models, config, image_processor, prompt, concepts, output_dir="concept_images"):
+    """
+    Generate images using multiple concepts
+    """
+    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
+    
+    for concept in concepts:
+        # ... your existing code ...
+        # Make sure to return the PIL Image object
+        return pil_image
+    
+    # If we get here, return None
+    return None
 
 def channel_loss(images, channel_idx=2, target_value=0.9):
     """