Update app.py

app.py

@@ -0,0 +1,311 @@
+import gradio as gr
+import torch
+from diffusers import StableDiffusionPipeline
+from PIL import Image
+import numpy as np
+import os
+from huggingface_hub import hf_hub_download
+import warnings
+from transformers import CLIPProcessor, CLIPModel
+warnings.filterwarnings("ignore")
+
+# Check if CUDA is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+# Load CLIP model for semantic guidance
+print("Loading CLIP model for semantic guidance...")
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+
+# Dictionary of available concepts
+CONCEPTS = {
+    "canna-lily-flowers102": {
+        "repo_id": "sd-concepts-library/canna-lily-flowers102",
+        "type": "object",
+        "description": "Canna lily flower style"
+    },
+    "samurai-jack": {
+        "repo_id": "sd-concepts-library/samurai-jack",
+        "type": "style",
+        "description": "Samurai Jack animation style"
+    },
+    "babies-poster": {
+        "repo_id": "sd-concepts-library/babies-poster",
+        "type": "style",
+        "description": "Babies poster art style"
+    },
+    "animal-toy": {
+        "repo_id": "sd-concepts-library/animal-toy",
+        "type": "object",
+        "description": "Animal toy style"
+    },
+    "sword-lily-flowers102": {
+        "repo_id": "sd-concepts-library/sword-lily-flowers102",
+        "type": "object",
+        "description": "Sword lily flower style"
+    }
+}
+
+def car_loss(image):
+    """Custom loss function that encourages the presence of cars in the image"""
+    # Convert PIL image to tensor if needed
+    if isinstance(image, Image.Image):
+        image = np.array(image)
+        image = torch.tensor(image, device=device)
+    
+    # Process image for CLIP
+    with torch.no_grad():
+        # Convert to PIL for CLIP processing
+        pil_image = Image.fromarray(image.cpu().numpy().astype(np.uint8))
+        
+        # Get CLIP features for the image
+        inputs = clip_processor(
+            text=["a photo of a car", "a photo without cars"],
+            images=pil_image,
+            return_tensors="pt",
+            padding=True
+        ).to(device)
+        
+        # Get similarity scores
+        outputs = clip_model(**inputs)
+        logits_per_image = outputs.logits_per_image
+        
+        # Higher score for the first text (with cars) is better
+        car_score = logits_per_image[0][0]
+        no_car_score = logits_per_image[0][1]
+        
+        # We want to maximize car_score and minimize no_car_score
+        loss = -(car_score - no_car_score)
+        
+    return loss
+
+def generate_image(pipe, prompt, seed, guidance_scale=7.5, num_inference_steps=30, use_car_guidance=False):
+    """Generate an image with optional car guidance"""
+    generator = torch.Generator(device).manual_seed(seed)
+    custom_loss = car_loss if use_car_guidance else None
+    
+    if custom_loss:
+        try:
+            # Start with a standard generation
+            init_images = pipe(
+                prompt,
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps // 2,
+                generator=generator
+            ).images
+            
+            init_image = init_images[0]
+            
+            # Refine using car guidance
+            from diffusers import StableDiffusionImg2ImgPipeline
+            
+            img2img_pipe = StableDiffusionImg2ImgPipeline(
+                vae=pipe.vae,
+                text_encoder=pipe.text_encoder,
+                tokenizer=pipe.tokenizer,
+                unet=pipe.unet,
+                scheduler=pipe.scheduler,
+                safety_checker=None,
+                feature_extractor=None,
+            ).to(device)
+            
+            strength = 0.75
+            current_image = init_image
+            
+            for i in range(5):
+                current_loss = custom_loss(current_image)
+                
+                refined_images = img2img_pipe(
+                    prompt=prompt + ", with beautiful cars",
+                    image=current_image,
+                    strength=strength,
+                    guidance_scale=guidance_scale,
+                    generator=generator,
+                ).images
+                
+                current_image = refined_images[0]
+                strength *= 0.8
+            
+            return current_image
+            
+        except Exception as e:
+            print(f"Error in car-guided generation: {e}")
+            return pipe(
+                prompt,
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps,
+                generator=generator
+            ).images[0]
+    else:
+        return pipe(
+            prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            generator=generator
+        ).images[0]
+
+# Cache for loaded models and concepts
+loaded_models = {}
+
+def get_model_with_concept(concept_name):
+    """Get or load a model with the specified concept"""
+    if concept_name not in loaded_models:
+        concept_info = CONCEPTS[concept_name]
+        
+        # Download concept embedding
+        concept_path = f"concepts/{concept_name}.bin"
+        os.makedirs("concepts", exist_ok=True)
+        
+        if not os.path.exists(concept_path):
+            file = hf_hub_download(
+                repo_id=concept_info["repo_id"],
+                filename="learned_embeds.bin",
+                repo_type="model"
+            )
+            import shutil
+            shutil.copy(file, concept_path)
+        
+        # Load model and concept
+        pipe = StableDiffusionPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-2",
+            torch_dtype=torch.float32 if device == "cpu" else torch.float16,
+            safety_checker=None
+        ).to(device)
+        
+        pipe.load_textual_inversion(concept_path)
+        loaded_models[concept_name] = pipe
+    
+    return loaded_models[concept_name]
+
+def generate_images(concept_name, base_prompt, seed, use_car_guidance):
+    """Generate images using the selected concept"""
+    try:
+        # Get model with concept
+        pipe = get_model_with_concept(concept_name)
+        
+        # Construct prompt based on concept type
+        if CONCEPTS[concept_name]["type"] == "object":
+            prompt = f"A {base_prompt} with a <{concept_name}>"
+        else:
+            prompt = f"<{concept_name}> {base_prompt}"
+        
+        # Generate image
+        image = generate_image(
+            pipe=pipe,
+            prompt=prompt,
+            seed=int(seed),
+            use_car_guidance=use_car_guidance
+        )
+        
+        return image
+    except Exception as e:
+        raise gr.Error(f"Error generating image: {str(e)}")
+
+# Create Gradio interface
+with gr.Blocks(title="Stable Diffusion Style Explorer") as demo:
+    gr.Markdown("""
+    # Stable Diffusion Style Explorer
+    
+    Generate images using various concepts from the SD Concepts Library, with optional car guidance.
+    
+    ## How to use:
+    1. Select a concept from the dropdown
+    2. Enter a base prompt (or use the default)
+    3. Set a seed for reproducibility
+    4. Choose whether to use car guidance
+    5. Click Generate!
+    
+    Check out the examples below to see different combinations of concepts and prompts!
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            concept = gr.Dropdown(
+                choices=list(CONCEPTS.keys()),
+                value="samurai-jack",
+                label="Select Concept"
+            )
+            
+            prompt = gr.Textbox(
+                value="A serene landscape with mountains and a lake at sunset",
+                label="Base Prompt"
+            )
+            
+            seed = gr.Number(
+                value=42,
+                label="Seed",
+                precision=0
+            )
+            
+            car_guidance = gr.Checkbox(
+                value=False,
+                label="Use Car Guidance"
+            )
+            
+            generate_btn = gr.Button("Generate Image")
+        
+        with gr.Column():
+            output_image = gr.Image(label="Generated Image")
+            
+    concept.change(
+        fn=lambda x: gr.Markdown(f"Selected concept: {CONCEPTS[x]['description']} ({CONCEPTS[x]['type']})"),
+        inputs=[concept],
+        outputs=[gr.Markdown()]
+    )
+    
+    generate_btn.click(
+        fn=generate_images,
+        inputs=[concept, prompt, seed, car_guidance],
+        outputs=[output_image]
+    )
+
+    # Gallery of pre-generated examples
+    gr.Markdown("### 🖼️ Pre-generated Examples")
+    
+    with gr.Row():
+        # Samurai Jack examples
+        with gr.Column():
+            gr.Markdown("**Samurai Jack Style**")
+            gr.Image("Assignment17/Assignment17/outputs/samurai-jack_normal.png", 
+                    label="Without Car Guidance")
+            gr.Image("Assignment17/Assignment17/outputs/samurai-jack_car.png", 
+                    label="With Car Guidance")
+    
+    with gr.Row():
+        # Canna Lily examples
+        with gr.Column():
+            gr.Markdown("**Canna Lily Object**")
+            gr.Image("Assignment17/Assignment17/outputs/canna-lily-flowers102_normal.png", 
+                    label="Without Car Guidance")
+            gr.Image("Assignment17/Assignment17/outputs/canna-lily-flowers102_car.png", 
+                    label="With Car Guidance")
+    
+    with gr.Row():
+        # Babies Poster examples
+        with gr.Column():
+            gr.Markdown("**Babies Poster Style**")
+            gr.Image("Assignment17/Assignment17/outputs/babies-poster_normal.png", 
+                    label="Without Car Guidance")
+            gr.Image("Assignment17/Assignment17/outputs/babies-poster_car.png", 
+                    label="With Car Guidance")
+    
+    with gr.Row():
+        # Animal Toy examples
+        with gr.Column():
+            gr.Markdown("**Animal Toy Object**")
+            gr.Image("Assignment17/Assignment17/outputs/animal-toy_normal.png", 
+                    label="Without Car Guidance")
+            gr.Image("Assignment17/Assignment17/outputs/animal-toy_car.png", 
+                    label="With Car Guidance")
+    
+    with gr.Row():
+        # Sword Lily examples
+        with gr.Column():
+            gr.Markdown("**Sword Lily Object**")
+            gr.Image("Assignment17/Assignment17/outputs/sword-lily-flowers102_normal.png", 
+                    label="Without Car Guidance")
+            gr.Image("Assignment17/Assignment17/outputs/sword-lily-flowers102_car.png", 
+                    label="With Car Guidance")
+
+demo.launch()