Upload 2 files

README.md

@@ -1,13 +1,129 @@
----
-title: Text Inversion
-emoji: 🖼
-colorFrom: purple
-colorTo: red
-sdk: gradio
-sdk_version: 5.0.1
-app_file: app.py
-pinned: false
-short_description: Stable Diffusion - Text Inversion
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Stable Diffusion Text Inversion with Loss Functions
+
+This repository contains a Gradio web application that provides an intuitive interface for generating images using Stable Diffusion with textual inversion and guided loss functions.
+
+## Overview
+
+The application allows users to explore the capabilities of Stable Diffusion by:
+- Generating images from text prompts
+- Using textual inversion concepts
+- Applying various loss functions to guide the diffusion process
+- Generating multiple images with different seeds
+
+!![alt text](image.png)
+
+## Features
+
+### Core Functionality
+- **Text-to-Image Generation**: Create detailed images from descriptive text prompts
+- **Textual Inversion**: Apply learned concepts to your generations
+- **Loss Function Guidance**: Shape image generation with specialized loss functions:
+  - **Blue Loss**: Emphasizes blue tones in the generated images
+  - **Elastic Loss**: Creates distortion effects by applying elastic transformations
+  - **Symmetry Loss**: Encourages symmetrical image generation
+  - **Saturation Loss**: Enhances color saturation in the output
+- **Multi-Seed Generation**: Create multiple variations of an image with different seeds
+
+## Installation
+
+### Prerequisites
+- Python 3.8+
+- CUDA-capable GPU (recommended)
+
+### Setup
+1. Clone this repository:
+```bash
+git clone https://github.com/yourusername/stable-diffusion-text-inversion.git
+cd stable-diffusion-text-inversion
+```
+
+2. Install dependencies:
+```bash
+pip install torch diffusers transformers tqdm torchvision matplotlib gradio
+```
+
+3. Run the application:
+```bash
+python gradio_app.py
+```
+
+4. Open the provided URL (typically http://localhost:7860) in your browser.
+
+## Understanding the Technology
+
+### Stable Diffusion
+
+Stable Diffusion is a latent text-to-image diffusion model developed by Stability AI. It works by:
+
+1. **Encoding text**: Converting text prompts into embeddings that the model can understand
+2. **Starting with noise**: Beginning with random noise in a latent space
+3. **Iterative denoising**: Gradually removing noise while being guided by the text embeddings
+4. **Decoding to image**: Converting the final latent representation to a pixel-based image
+
+The model operates in a compressed latent space (64x64x4) rather than pixel space (512x512x3), allowing for efficient generation of high-resolution images with limited computational resources.
+
+### Textual Inversion
+
+Textual Inversion is a technique that allows Stable Diffusion to learn new concepts from just a few example images. Key aspects include:
+
+- **Custom Concepts**: Learn new visual concepts not present in the model's training data
+- **Few-Shot Learning**: Typically requires only 3-5 examples of a concept
+- **Token Optimization**: Creates a new "pseudo-word" embedding that represents the concept
+- **Seamless Integration**: Once learned, concepts can be used in prompts just like regular words
+
+In this application, we load several pre-trained textual inversion concepts from the SD concepts library:
+- Rimworld art style
+- HK Golden Lantern
+- Phoenix-01
+- Fractal Flame
+- Scarlet Witch
+
+### Guided Loss Functions
+
+This application introduces an innovative approach by applying custom loss functions during the diffusion process:
+
+1. **How it works**: During generation, we periodically decode the current latent representation, apply a loss function to the decoded image, and backpropagate that loss to adjust the latents.
+
+2. **Types of Loss Functions**:
+   - **Blue Loss**: Encourages pixels to have higher values in the blue channel
+   - **Elastic Loss**: Minimizes difference between the image and an elastically transformed version
+   - **Symmetry Loss**: Minimizes difference between the image and its horizontal mirror
+   - **Saturation Loss**: Pushes the image toward higher color saturation
+
+3. **Impact**: These loss functions can dramatically alter the aesthetic qualities of the generated images, creating effects that would be difficult to achieve through prompt engineering alone.
+
+## Usage Examples
+
+### Basic Image Generation
+1. Enter a prompt in the text box (e.g., "A majestic castle on a floating island with waterfalls")
+2. Set Loss Type to "N/A" and uncheck "Apply Loss Function"
+3. Enter a seed value (e.g., "42")
+4. Click "Generate Images"
+
+### Applying Loss Functions
+1. Enter your prompt
+2. Select a Loss Type (e.g., "symmetry")
+3. Check "Apply Loss Function"
+4. Enter a seed value
+5. Click "Generate Images"
+
+### Batch Generation
+1. Enter your prompt
+2. Select desired loss settings
+3. Enter multiple comma-separated seeds (e.g., "42, 100, 500")
+4. Click "Generate Images" to generate a grid of variations
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request.
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details.
+
+## Acknowledgments
+
+- [Stability AI](https://stability.ai/) for developing Stable Diffusion
+- [Hugging Face](https://huggingface.co/) for the Diffusers library
+- [Gradio](https://gradio.app/) for the web interface framework
+- The creators of the textual inversion concepts used in this project

gradio_app.py

@@ -0,0 +1,279 @@
+import os
+import torch
+import gradio as gr
+from PIL import Image
+from diffusers import StableDiffusionPipeline, DiffusionPipeline
+from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel
+from tqdm.auto import tqdm
+import torchvision.transforms as T
+import torch.nn.functional as F
+import gc
+
+# Configure constants
+HEIGHT, WIDTH = 512, 512
+GUIDANCE_SCALE = 8
+LOSS_SCALE = 200
+NUM_INFERENCE_STEPS = 50
+BATCH_SIZE = 1
+DEFAULT_PROMPT = "A deadly witcher slinging a sword with a lion medallion in his neck, casting a fire spell from his hand in a snowy forest"
+
+# Define the device
+TORCH_DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+
+# Initialize the elastic transformer
+elastic_transformer = T.ElasticTransform(alpha=550.0, sigma=5.0)
+
+# Load the model
+def load_model():
+    pipe = DiffusionPipeline.from_pretrained(
+        "CompVis/stable-diffusion-v1-4",
+        torch_dtype=torch.float16 if TORCH_DEVICE == "cuda" else torch.float32
+    ).to(TORCH_DEVICE)
+    
+    # Load textual inversion concepts
+    try:
+        pipe.load_textual_inversion("sd-concepts-library/rimworld-art-style", mean_resizing=False)
+        pipe.load_textual_inversion("sd-concepts-library/hk-goldenlantern", mean_resizing=False)
+        pipe.load_textual_inversion("sd-concepts-library/phoenix-01", mean_resizing=False)
+        pipe.load_textual_inversion("sd-concepts-library/fractal-flame", mean_resizing=False)
+        pipe.load_textual_inversion("sd-concepts-library/scarlet-witch", mean_resizing=False)
+    except Exception as e:
+        print(f"Warning: Could not load all textual inversion concepts: {e}")
+        
+    return pipe
+
+# Helper functions
+def image_grid(imgs, rows, cols):
+    assert len(imgs) == rows*cols
+    w, h = imgs[0].size
+    grid = Image.new('RGB', size=(cols*w, rows*h))
+    
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i%cols*w, i//cols*h))
+    return grid
+
+def image_loss(images, loss_type):
+    if loss_type == 'blue':
+        # blue loss
+        error = torch.abs(images[:,2] - 0.9).mean()
+    elif loss_type == 'elastic':
+        # elastic loss
+        transformed_imgs = elastic_transformer(images)
+        error = torch.abs(transformed_imgs - images).mean()
+    elif loss_type == 'symmetry':
+        flipped_image = torch.flip(images, [3])
+        error = F.mse_loss(images, flipped_image)
+    elif loss_type == 'saturation':
+        # saturation loss
+        transformed_imgs = T.functional.adjust_saturation(images, saturation_factor=10)
+        error = torch.abs(transformed_imgs - images).mean()
+    else:
+        print("Error. Loss not defined")
+        error = torch.tensor(0.0)
+
+    return error
+
+def latents_to_pil(latents, pipe):
+    # batch of latents -> list of images
+    latents = (1 / 0.18215) * latents
+    with torch.no_grad():
+        image = pipe.vae.decode(latents).sample
+    image = (image / 2 + 0.5).clamp(0, 1)
+    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
+    images = (image * 255).round().astype("uint8")
+    pil_images = [Image.fromarray(image) for image in images]
+    return pil_images
+
+def generate_image(pipe, seed_no, prompts, loss_type, loss_apply=False, progress=gr.Progress()):
+    # Initialization and Setup
+    generator = torch.manual_seed(seed_no)
+
+    scheduler = LMSDiscreteScheduler(
+        beta_start=0.00085, 
+        beta_end=0.012, 
+        beta_schedule="scaled_linear", 
+        num_train_timesteps=1000
+    )
+    scheduler.set_timesteps(NUM_INFERENCE_STEPS)
+    scheduler.timesteps = scheduler.timesteps.to(torch.float32)
+
+    # Text Processing
+    text_input = pipe.tokenizer(
+        prompts, 
+        padding='max_length', 
+        max_length=pipe.tokenizer.model_max_length, 
+        truncation=True, 
+        return_tensors="pt"
+    )
+    input_ids = text_input.input_ids.to(TORCH_DEVICE)
+
+    # Convert text inputs to embeddings
+    with torch.no_grad():
+        text_embeddings = pipe.text_encoder(input_ids)[0]
+
+    # Handle padding and truncation of text inputs
+    max_length = text_input.input_ids.shape[-1]
+    uncond_input = pipe.tokenizer(
+        [""] * BATCH_SIZE, 
+        padding="max_length", 
+        max_length=max_length, 
+        return_tensors="pt"
+    )
+
+    with torch.no_grad():
+        uncond_embeddings = pipe.text_encoder(uncond_input.input_ids.to(TORCH_DEVICE))[0]
+
+    # Concatenate unconditioned and text embeddings
+    text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
+
+    # Create random initial latents
+    latents = torch.randn(
+        (BATCH_SIZE, pipe.unet.config.in_channels, HEIGHT // 8, WIDTH // 8),
+        generator=generator,
+    )
+
+    # Move latents to device and apply noise scaling
+    if TORCH_DEVICE == "cuda":
+        latents = latents.to(torch.float16)
+    latents = latents.to(TORCH_DEVICE)
+    latents = latents * scheduler.init_noise_sigma
+
+    # Diffusion Process
+    for i, t in progress.tqdm(enumerate(scheduler.timesteps), total=len(scheduler.timesteps)):
+        # Process the latent model input
+        latent_model_input = torch.cat([latents] * 2)
+        sigma = scheduler.sigmas[i]
+        latent_model_input = scheduler.scale_model_input(latent_model_input, t)
+
+        with torch.no_grad():
+            noise_pred = pipe.unet(
+                latent_model_input, 
+                t, 
+                encoder_hidden_states=text_embeddings
+            )["sample"]
+
+        # Apply noise prediction
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + GUIDANCE_SCALE * (noise_pred_text - noise_pred_uncond)
+
+        # Apply loss if requested
+        if loss_apply and i % 5 == 0:
+            latents = latents.detach().requires_grad_()
+            latents_x0 = latents - sigma * noise_pred
+
+            # Use VAE to decode the image
+            denoised_images = pipe.vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5
+
+            # Apply loss
+            loss = image_loss(denoised_images, loss_type) * LOSS_SCALE
+            print(f"Step {i}, Loss: {loss.item()}")
+
+            # Compute gradients for optimization
+            cond_grad = torch.autograd.grad(loss, latents)[0]
+            latents = latents.detach() - cond_grad * sigma**2
+
+        # Update latents using the scheduler
+        latents = scheduler.step(noise_pred, t, latents).prev_sample
+
+    return latents
+
+def generate_images(prompt, loss_type, apply_loss, seeds, pipe):
+    latents_collect = []
+    
+    # Convert comma-separated string to list and clean
+    seeds = [int(seed.strip()) for seed in seeds.split(',') if seed.strip()]
+    
+    if not seeds:
+        seeds = [1000]  # Default seed if none provided
+        
+    # List of SD concepts (can be empty if not used)
+    sdconcepts = [''] * len(seeds)
+    
+    # Generate images for each seed
+    for seed_no, sd in zip(seeds, sdconcepts):
+        # Clear CUDA cache
+        if TORCH_DEVICE == "cuda":
+            torch.cuda.empty_cache()
+            gc.collect()
+            torch.cuda.empty_cache()
+        
+        # Generate image
+        prompts = [f'{prompt} {sd}']
+        latents = generate_image(pipe, seed_no, prompts, loss_type, loss_apply=apply_loss)
+        latents_collect.append(latents)
+    
+    # Stack latents and convert to images
+    latents_collect = torch.vstack(latents_collect)
+    images = latents_to_pil(latents_collect, pipe)
+    
+    # Create image grid
+    if len(images) > 1:
+        result = image_grid(images, 1, len(images))
+        return result
+    else:
+        return images[0]
+
+# Gradio Interface
+def create_interface():
+    pipe = load_model()
+    
+    with gr.Blocks(title="Stable Diffusion Text Inversion with Loss Functions") as app:
+        gr.Markdown("""
+        # Stable Diffusion Text Inversion with Loss Functions
+        
+        Generate images using Stable Diffusion with various loss functions to guide the diffusion process.
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                prompt = gr.Textbox(
+                    label="Prompt", 
+                    value=DEFAULT_PROMPT,
+                    lines=3
+                )
+                
+                loss_type = gr.Radio(
+                    label="Loss Type",
+                    choices=["N/A", "blue", "elastic", "symmetry", "saturation"],
+                    value="N/A"
+                )
+                
+                apply_loss = gr.Checkbox(
+                    label="Apply Loss Function", 
+                    value=False
+                )
+                
+                seeds = gr.Textbox(
+                    label="Seeds (comma-separated)",
+                    value="3000,2000,1000",
+                    lines=1
+                )
+                
+                generate_btn = gr.Button("Generate Images")
+                
+            with gr.Column():
+                output_image = gr.Image(label="Generated Image")
+                
+        generate_btn.click(
+            fn=lambda p, lt, al, s: generate_images(p, lt, al, s, pipe),
+            inputs=[prompt, loss_type, apply_loss, seeds],
+            outputs=output_image
+        )
+        
+        gr.Markdown("""
+        ## About the Loss Functions
+        
+        - **Blue**: Encourages more blue tones in the image
+        - **Elastic**: Creates distortion effects by minimizing differences with elastically transformed versions
+        - **Symmetry**: Encourages symmetrical images by minimizing differences with horizontally flipped versions
+        - **Saturation**: Increases color saturation in the image
+        
+        Set "N/A" and uncheck "Apply Loss Function" for normal image generation.
+        """)
+    
+    return app
+
+if __name__ == "__main__":
+    # Create and launch the interface
+    app = create_interface()
+    app.launch()