Create app.py

initial app commit

app.py

@@ -0,0 +1,217 @@
+import os
+import json
+import torch
+import gc
+import numpy as np
+import gradio as gr
+from PIL import Image
+from diffusers import StableDiffusionXLPipeline
+import open_clip
+from huggingface_hub import hf_hub_download
+from IP_Adapter.ip_adapter import IPAdapterXL
+from perform_swap import compute_dataset_embeds_svd, get_modified_images_embeds_composition
+import tempfile
+import uuid
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Initialize SDXL pipeline
+base_model_path = "stabilityai/stable-diffusion-xl-base-1.0"
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    base_model_path,
+    torch_dtype=torch.float16,
+    add_watermarker=False,
+)
+
+# Initialize IP-Adapter
+image_encoder_repo = 'h94/IP-Adapter'
+image_encoder_subfolder = 'models/image_encoder'
+ip_ckpt = hf_hub_download('h94/IP-Adapter', subfolder="sdxl_models", filename='ip-adapter_sdxl_vit-h.bin')
+ip_model = IPAdapterXL(pipe, image_encoder_repo, image_encoder_subfolder, ip_ckpt, device)
+
+# Initialize CLIP model
+clip_model, _, preprocess = open_clip.create_model_and_transforms('hf-hub:laion/CLIP-ViT-H-14-laion2B-s32B-b79K')
+clip_model.to(device)
+print("Models initialized successfully!")
+
+def get_image_embeds(pil_image, model=clip_model, preproc=preprocess, dev=device):
+    """Get CLIP image embeddings for a given PIL image"""
+    image = preproc(pil_image)[np.newaxis, :, :, :]
+    with torch.no_grad():
+        embeds = model.encode_image(image.to(dev))
+    return embeds.cpu().detach().numpy()
+
+def save_temp_image(image):
+    """Save a PIL image to a temporary file and return the path"""
+    temp_dir = tempfile.gettempdir()
+    filename = f"{uuid.uuid4()}.png"
+    filepath = os.path.join(temp_dir, filename)
+    image.save(filepath)
+    return filepath
+
+def process_images(
+    base_image, 
+    concept_image1, concept_desc1,
+    concept_image2=None, concept_desc2=None,
+    concept_image3=None, concept_desc3=None,
+    rank1=10, rank2=10, rank3=10,
+    prompt=None, 
+    scale=1.0,
+    seed=420
+):
+    """Process the base image and concept images to generate modified images"""
+    # Process base image
+    base_image_pil = Image.fromarray(base_image).convert("RGB")
+    base_embed = get_image_embeds(base_image_pil)
+    
+    # Process concept images
+    concept_images = []
+    concept_descriptions = []
+    
+    # Add first concept (required)
+    if concept_image1 is not None:
+        concept_images.append(concept_image1)
+        concept_descriptions.append(concept_desc1 if concept_desc1 else "Concept 1")
+    else:
+        return None, "Please upload at least one concept image"
+    
+    # Add second concept (optional)
+    if concept_image2 is not None:
+        concept_images.append(concept_image2)
+        concept_descriptions.append(concept_desc2 if concept_desc2 else "Concept 2")
+    
+    # Add third concept (optional)
+    if concept_image3 is not None:
+        concept_images.append(concept_image3)
+        concept_descriptions.append(concept_desc3 if concept_desc3 else "Concept 3")
+    
+    # Get all ranks
+    ranks = [rank1]
+    if concept_image2 is not None:
+        ranks.append(rank2)
+    if concept_image3 is not None:
+        ranks.append(rank3)
+    
+    concept_embeds = []
+    for img in concept_images:
+        if img is not None:
+            img_pil = Image.fromarray(img).convert("RGB")
+            concept_embeds.append(get_image_embeds(img_pil))
+    
+    # Compute projection matrices
+    projection_matrices = []
+    for i, embed in enumerate(concept_embeds):
+        # For a single image, we need to reshape to have the same format as a collection
+        single_embed = embed.reshape(1, *embed.shape)
+        projection_matrix = compute_dataset_embeds_svd(single_embed, ranks[i])
+        projection_matrices.append(projection_matrix)
+    
+    # Create projection data structure for the composition
+    projections_data = [
+        {
+            "embed": embed,
+            "projection_matrix": proj_matrix
+        }
+        for embed, proj_matrix in zip(concept_embeds, projection_matrices)
+    ]
+    
+    # Generate modified images - 
+    modified_images = get_modified_images_embeds_composition(
+        base_embed, 
+        projections_data, 
+        ip_model, 
+        prompt=prompt, 
+        scale=scale, 
+        num_samples=1, 
+        seed=seed
+    )
+    
+    return modified_images
+
+def process_and_display(
+    base_image, 
+    concept_image1, concept_desc1,
+    concept_image2=None, concept_desc2=None,
+    concept_image3=None, concept_desc3=None,
+    rank1=10, rank2=10, rank3=10,
+    prompt=None, scale=1.0, seed=420
+):
+    """Wrapper for process_images that handles UI updates"""
+    if base_image is None:
+        return None, "Please upload a base image"
+    
+    if concept_image1 is None:
+        return None, "Please upload at least one concept image"
+    
+    modified_images = process_images(
+        base_image, 
+        concept_image1, concept_desc1,
+        concept_image2, concept_desc2,
+        concept_image3, concept_desc3,
+        rank1, rank2, rank3, 
+        prompt, scale, seed
+    )
+    
+    # # Clean up memory
+    # torch.cuda.empty_cache()
+    # gc.collect()
+    
+    return modified_images
+
+with gr.Blocks(title="Image Concept Composition") as demo:
+        gr.Markdown("# Image Concept Composition")
+        gr.Markdown("Upload a base image and 1-3 concept images to create new images that combine these concepts.")
+        
+        with gr.Row():
+            with gr.Column():
+                base_image = gr.Image(label="Base Image (Required)", type="numpy")
+                
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        concept_image1 = gr.Image(label="Concept Image 1 (Required)", type="numpy")
+                    with gr.Column(scale=1):
+                        concept_desc1 = gr.Textbox(label="Concept 1 Description", placeholder="Describe this concept")
+                        rank1 = gr.Slider(minimum=1, maximum=50, value=10, step=1, label="Rank 1")
+                
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        concept_image2 = gr.Image(label="Concept Image 2 (Optional)", type="numpy")
+                    with gr.Column(scale=1):
+                        concept_desc2 = gr.Textbox(label="Concept 2 Description", placeholder="Describe this concept")
+                        rank2 = gr.Slider(minimum=1, maximum=50, value=10, step=1, label="Rank 2")
+                
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        concept_image3 = gr.Image(label="Concept Image 3 (Optional)", type="numpy")
+                    with gr.Column(scale=1):
+                        concept_desc3 = gr.Textbox(label="Concept 3 Description", placeholder="Describe this concept")
+                        rank3 = gr.Slider(minimum=1, maximum=50, value=10, step=1, label="Rank 3")
+                
+                prompt = gr.Textbox(label="Guidance Prompt (Optional)", placeholder="Optional text prompt to guide generation")
+                
+                with gr.Row():
+                    scale = gr.Slider(minimum=0.1, maximum=2.0, value=1.0, step=0.1, label="Scale")
+                    seed = gr.Number(value=420, label="Seed", precision=0)
+                
+                submit_btn = gr.Button("Generate Image")
+            
+            with gr.Column():
+                gallery = gr.Gallery(label="Generated Image", show_label=True)
+                status = gr.Markdown("Upload images and click Generate")
+        
+        submit_btn.click(
+            fn=process_and_display,
+            inputs=[
+                base_image, 
+                concept_image1, concept_desc1,
+                concept_image2, concept_desc2,
+                concept_image3, concept_desc3,
+                rank1, rank2, rank3, 
+                prompt, scale, seed
+            ],
+            outputs=[gallery, status]
+        )
+    
+    
+
+    demo.launch()