Modified with Examples

Modified with Examples.

app.py

@@ -6,6 +6,8 @@ from huggingface_hub import hf_hub_download
 import os
 from pathlib import Path
 import traceback
+import glob
+from PIL import Image
 
 # Reuse the same load_learned_embed_in_clip and Distance_loss functions
 def load_learned_embed_in_clip(learned_embeds_path, text_encoder, tokenizer, token=None):
@@ -92,7 +94,6 @@ class StyleGenerator:
             "Bird Style"
         ]
         self.is_initialized = False
-        # Check if CUDA is available
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         if self.device == "cpu":
             print("NVIDIA GPU not found. Running on CPU (this will be slower)")
@@ -132,63 +133,44 @@ class StyleGenerator:
             print(traceback.format_exc())
             raise
 
-    def generate_images(self, prompt, apply_loss=False, num_inference_steps=50, guidance_scale=7.5):
-        if not self.is_initialized:
-            self.initialize_model()
-            
-        images = []
-        style_names = []
-        
+    def generate_single_style(self, prompt, selected_style):
         try:
-            def callback_fn(i, t, latents):
-                if i % 5 == 0 and apply_loss:
-                    try:
-                        # Ensure latents are in the correct format and require gradients
-                        latents = latents.float()
-                        latents.requires_grad_(True)
-                        
-                        # Compute loss
-                        loss = Distance_loss(latents)
-                        
-                        # Compute gradients manually
-                        grads = torch.autograd.grad(
-                            outputs=loss,
-                            inputs=latents,
-                            create_graph=False,
-                            retain_graph=False,
-                            only_inputs=True
-                        )[0]
-                        
-                        # Update latents
-                        with torch.no_grad():
-                            latents = latents - 0.1 * grads
-                            
-                    except Exception as e:
-                        print(f"Error in callback: {e}")
-                        return latents
-                        
-                return latents
-
-            for style_token, style_name in zip(self.style_tokens, self.style_names):
-                styled_prompt = f"{prompt}, {style_token}"
-                style_names.append(style_name)
-                
-                # Disable autocast for better gradient computation
-                image = self.pipe(
+            # Find the index of the selected style
+            style_idx = self.style_names.index(self.style_names[selected_style])
+            
+            # Generate single image with selected style
+            styled_prompt = f"{prompt}, {self.style_tokens[style_idx]}"
+            
+            # Set seed for reproducibility
+            generator_seed = 42
+            torch.manual_seed(generator_seed)
+            if self.device == "cuda":
+                torch.cuda.manual_seed(generator_seed)
+            
+            # Generate base image
+            with autocast(self.device):
+                base_image = self.pipe(
                     styled_prompt,
-                    num_inference_steps=num_inference_steps,
-                    guidance_scale=guidance_scale,
-                    callback=callback_fn if apply_loss else None,
-                    callback_steps=5
+                    num_inference_steps=50,
+                    guidance_scale=7.5,
+                    generator=torch.Generator(self.device).manual_seed(generator_seed)
                 ).images[0]
-                
-                images.append(image)
             
-            return images, style_names
+            # Generate same image with loss
+            with autocast(self.device):
+                loss_image = self.pipe(
+                    styled_prompt,
+                    num_inference_steps=50,
+                    guidance_scale=7.5,
+                    callback=self.callback_fn,
+                    callback_steps=5,
+                    generator=torch.Generator(self.device).manual_seed(generator_seed)
+                ).images[0]
+            
+            return base_image, loss_image
             
         except Exception as e:
-            print(f"Error during image generation: {str(e)}")
-            print(traceback.format_exc())
+            print(f"Error in generate_single_style: {e}")
             raise
 
     def callback_fn(self, i, t, latents):
@@ -219,38 +201,44 @@ class StyleGenerator:
             
         return latents
 
-def generate_all_variations(prompt):
+def generate_single_style(prompt, selected_style):
     try:
         generator = StyleGenerator.get_instance()
         if not generator.is_initialized:
             generator.initialize_model()
         
-        # Generate images without loss
-        regular_images, style_names = generator.generate_images(prompt, apply_loss=False)
+        base_image, loss_image = generator.generate_single_style(prompt, selected_style)
         
-        # Generate images with loss
-        loss_images, _ = generator.generate_images(prompt, apply_loss=True)
-        
-        return regular_images, loss_images, style_names
-    
+        return [
+            gr.update(visible=False),  # error_message
+            base_image,                # original_image
+            loss_image                 # loss_image
+        ]
     except Exception as e:
-        print(f"Error in generate_all_variations: {str(e)}")
-        print(traceback.format_exc())
-        raise
+        print(f"Error in generate_single_style: {e}")
+        return [
+            gr.update(value=f"Error: {str(e)}", visible=True),  # error_message
+            None,  # original_image
+            None   # loss_image
+        ]
 
-def gradio_interface(prompt):
-    try:
-        regular_images, loss_images, style_names = generate_all_variations(prompt)
-        
-        return (
-            regular_images,  # Just return the images directly
-            loss_images     # Just return the images directly
-        )
-    except Exception as e:
-        print(f"Error in interface: {str(e)}")
-        print(traceback.format_exc())
-        # Return empty lists in case of error
-        return [], []
+# Add at the start of your script
+def debug_image_paths():
+    output_dir = Path("Outputs")
+    enhanced_dir = output_dir / "Color_Enhanced"
+    print(f"\nChecking image paths:")
+    print(f"Current working directory: {Path.cwd()}")
+    print(f"Looking for images in: {enhanced_dir.absolute()}")
+    
+    if enhanced_dir.exists():
+        print("\nFound files:")
+        for file in enhanced_dir.glob("*.webp"):
+            print(f"- {file.name}")
+    else:
+        print("\nDirectory not found!")
+
+# Call this function before creating the interface
+debug_image_paths()
 
 # Create a more beautiful interface with custom styling
 with gr.Blocks(css="""
@@ -265,73 +253,43 @@ with gr.Blocks(css="""
         border: 1px solid #374151;
         color: #f3f4f6;
     }
+    /* Enhanced Tab Styling */
+    .tabs.svelte-710i53 {
+        margin-bottom: 0 !important;
+    }
+    .tab-nav.svelte-710i53 {
+        background: transparent !important;
+        border: none !important;
+        padding: 12px 24px !important;
+        margin: 0 2px !important;
+        color: #9CA3AF !important;
+        font-weight: 500 !important;
+        transition: all 0.2s ease !important;
+        border-bottom: 2px solid transparent !important;
+    }
+    .tab-nav.svelte-710i53.selected {
+        background: transparent !important;
+        color: #F3F4F6 !important;
+        border-bottom: 2px solid #6366F1 !important;
+    }
+    .tab-nav.svelte-710i53:hover {
+        color: #F3F4F6 !important;
+        border-bottom: 2px solid #4F46E5 !important;
+    }
 """) as iface:
-    # Header section with dark theme
+    # Header section
     gr.Markdown(
         """
-        <div class="dark-theme" style="text-align: center; max-width: 800px; margin: 0 auto;">
+        <div class="dark-theme" style="text-align: center;">
         # 🎨 AI Style Transfer Studio
-        ### Transform your ideas into artistic masterpieces with custom styles and enhanced colors
+        ### Transform your ideas into artistic masterpieces
         </div>
         """
     )
 
-    # Define the generate_single_style function first
-    def generate_single_style(prompt, selected_style):
-        try:
-            generator = StyleGenerator.get_instance()
-            if not generator.is_initialized:
-                generator.initialize_model()
-            
-            # Find the index of the selected style
-            style_idx = generator.style_names.index(generator.style_names[selected_style])
-            
-            # Generate single image with selected style
-            styled_prompt = f"{prompt}, {generator.style_tokens[style_idx]}"
-            
-            # Set seed for reproducibility
-            generator_seed = 42
-            torch.manual_seed(generator_seed)
-            if generator.device == "cuda":
-                torch.cuda.manual_seed(generator_seed)
-            
-            # Generate base image
-            with autocast(generator.device):
-                base_image = generator.pipe(
-                    styled_prompt,
-                    num_inference_steps=50,
-                    guidance_scale=7.5,
-                    generator=torch.Generator(generator.device).manual_seed(generator_seed)
-                ).images[0]
-            
-            # Generate same image with loss
-            with autocast(generator.device):
-                loss_image = generator.pipe(
-                    styled_prompt,
-                    num_inference_steps=50,
-                    guidance_scale=7.5,
-                    callback=generator.callback_fn,
-                    callback_steps=5,
-                    generator=torch.Generator(generator.device).manual_seed(generator_seed)
-                ).images[0]
-            
-            return [
-                gr.update(visible=False),  # error_message
-                base_image,                # original_image
-                loss_image                 # loss_image
-            ]
-        except Exception as e:
-            print(f"Error in generate_single_style: {e}")
-            return [
-                gr.update(value=f"Error: {str(e)}", visible=True),  # error_message
-                None,  # original_image
-                None   # loss_image
-            ]
-
-    # Main content
+    # Controls section
     with gr.Row():
-        # Left sidebar for controls
-        with gr.Column(scale=1, min_width=300):
+        with gr.Column(scale=1):
             gr.Markdown("## 🎯 Controls")
             
             prompt = gr.Textbox(
@@ -359,29 +317,88 @@ with gr.Blocks(css="""
                 size="lg"
             )
             
-            # Error messages
             error_message = gr.Markdown(visible=False)
-            
-            # Style description
             style_description = gr.Markdown()
-        
-        # Right side for image display
-        with gr.Column(scale=2):
-            gr.Markdown("## 🖼️ Generated Artwork")
-            with gr.Row():
-                with gr.Column():
-                    original_image = gr.Image(
-                        label="Original Style",
-                        show_label=True,
-                        height=400
-                    )
-                with gr.Column():
-                    loss_image = gr.Image(
-                        label="Color Enhanced",
-                        show_label=True,
-                        height=400
-                    )
-    
+
+    # Generated Images
+    with gr.Row():
+        with gr.Column():
+            original_image = gr.Image(
+                label="Original Style",
+                show_label=True,
+                height=300
+            )
+        with gr.Column():
+            loss_image = gr.Image(
+                label="Color Enhanced",
+                show_label=True,
+                height=300
+            )
+
+    # Example Gallery
+    gr.Markdown(
+        """
+        <div class="dark-theme">
+        ## 🎆 Example Gallery
+        Compare original and enhanced versions for each style:
+        </div>
+        """
+    )
+
+    # Example Images
+    with gr.Row():
+        try:
+            output_dir = Path("Outputs")
+            original_dir = output_dir
+            enhanced_dir = output_dir / "Color_Enhanced"
+
+            if enhanced_dir.exists():
+                original_images = {
+                    Path(f).stem.split('_example')[0]: f 
+                    for f in original_dir.glob("*.webp") 
+                    if '_example' in f.name
+                }
+                enhanced_images = {
+                    Path(f).stem.split('_example')[0]: f 
+                    for f in enhanced_dir.glob("*.webp") 
+                    if '_example' in f.name
+                }
+
+                styles = [
+                    ("ronaldo", "Ronaldo Style"),
+                    ("canna_lily", "Canna Lily"),
+                    ("three_stooges", "Three Stooges"),
+                    ("pop_art", "Pop Art"),
+                    ("bird_style", "Bird Style")
+                ]
+
+                # Create a grid of all styles
+                for style_key, style_name in styles:
+                    if style_key in original_images and style_key in enhanced_images:
+                        with gr.Row():
+                            gr.Markdown(f"### {style_name}")
+                        with gr.Row():
+                            with gr.Column(scale=1):
+                                gr.Image(
+                                    value=str(original_images[style_key]),
+                                    label="Original",
+                                    show_label=True,
+                                    height=180
+                                )
+                            with gr.Column(scale=1):
+                                gr.Image(
+                                    value=str(enhanced_images[style_key]),
+                                    label="Color Enhanced",
+                                    show_label=True,
+                                    height=180
+                                )
+                        # Add a small spacing between styles
+                        gr.Markdown("<div style='margin: 10px 0;'></div>")
+
+        except Exception as e:
+            print(f"Error in example gallery: {e}")
+            gr.Markdown(f"Error loading example gallery: {str(e)}")
+
     # Info section
     with gr.Row():
         with gr.Column():
@@ -441,7 +458,7 @@ with gr.Blocks(css="""
             "Specialized in capturing the beauty of nature and wildlife"
         ]
         styles = ["Ronaldo Style", "Canna Lily", "Three Stooges", "Pop Art", "Bird Style"]
-        return f"### Selected: {styles[style_idx]}\n{descriptions[style_idx]}"
+        return f"### Selected Style: {styles[style_idx]}\n{descriptions[style_idx]}"
 
     style_radio.change(
         fn=update_style_description,
@@ -449,7 +466,6 @@ with gr.Blocks(css="""
         outputs=style_description
     )
 
-    # Connect the generate button
     generate_btn.click(
         fn=generate_single_style,
         inputs=[prompt, style_radio],