cpu_optimize

app.py

@@ -8,39 +8,67 @@ from tqdm.auto import tqdm
 import torchvision.transforms as T
 import torch.nn.functional as F
 import gc
+import signal
+import time
+import traceback
 
-# Configure constants
-HEIGHT, WIDTH = 512, 512
-GUIDANCE_SCALE = 8
+# Configure constants - optimized for CPU
+HEIGHT, WIDTH = 384, 384  # Smaller images use less memory
+GUIDANCE_SCALE = 7.5
 LOSS_SCALE = 200
-NUM_INFERENCE_STEPS = 50
+NUM_INFERENCE_STEPS = 30  # Reduced from 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"
+print(f"Using device: {TORCH_DEVICE}")
 
 # Initialize the elastic transformer
 elastic_transformer = T.ElasticTransform(alpha=550.0, sigma=5.0)
 
+# Timeout handler for CPU processing
+def timeout_handler(signum, frame):
+    raise TimeoutError("Image generation took too long")
+
 # 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}")
+        pipe = DiffusionPipeline.from_pretrained(
+            "CompVis/stable-diffusion-v1-4",
+            torch_dtype=torch.float16 if TORCH_DEVICE == "cuda" else torch.float32,
+            safety_checker=None,  # Disable safety checker for memory
+            low_cpu_mem_usage=True  # Enable memory optimization
+        ).to(TORCH_DEVICE)
         
-    return pipe
+        # Load textual inversion for all devices including CPU
+        try:
+            # Load one at a time with memory cleanup between each
+            concepts = [
+                "sd-concepts-library/rimworld-art-style",
+                "sd-concepts-library/hk-goldenlantern",
+                "sd-concepts-library/phoenix-01",
+                "sd-concepts-library/fractal-flame",
+                "sd-concepts-library/scarlet-witch"
+            ]
+            
+            for concept in concepts:
+                try:
+                    print(f"Loading textual inversion concept: {concept}")
+                    pipe.load_textual_inversion(concept, mean_resizing=False)
+                    # Clear memory after loading each concept
+                    if TORCH_DEVICE == "cpu":
+                        gc.collect()
+                except Exception as e:
+                    print(f"Warning: Could not load textual inversion concept {concept}: {e}")
+        except Exception as e:
+            print(f"Warning: Could not load textual inversion concepts: {e}")
+            
+        return pipe
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        traceback.print_exc()
+        raise
 
 # Helper functions
 def image_grid(imgs, rows, cols):
@@ -85,145 +113,178 @@ def latents_to_pil(latents, pipe):
     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)
+    try:
+        # Set timeout for CPU
+        if TORCH_DEVICE == "cpu":
+            signal.signal(signal.SIGALRM, timeout_handler)
+            signal.alarm(600)  # 10 minute timeout
+            
+        # 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)
+        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)
+        # 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]
+        # 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"
-    )
+        # 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]
+        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])
+        # 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,
-    )
+        # 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
+        # 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
-    timesteps = scheduler.timesteps
-    progress(0, desc="Generating")
-    
-    # Fixed loop - separate the progress tracking from the enumeration
-    for i in range(len(timesteps)):
-        progress((i + 1) / len(timesteps), desc=f"Diffusion step {i+1}/{len(timesteps)}")
-        t = timesteps[i]
+        # Diffusion Process
+        timesteps = scheduler.timesteps
+        progress(0, desc="Generating")
         
-        # 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)
+        # Fixed loop - separate the progress tracking from the enumeration
+        for i in range(len(timesteps)):
+            progress((i + 1) / len(timesteps), desc=f"Diffusion step {i+1}/{len(timesteps)}")
+            t = timesteps[i]
+            
+            # 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"]
+            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 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
+            # Apply loss if requested
+            if loss_apply and i % 5 == 0 and loss_type != "N/A":
+                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
+                # 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()}")
+                # 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
+                # 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
+            # Update latents using the scheduler
+            latents = scheduler.step(noise_pred, t, latents).prev_sample
+            
+            # Garbage collect every 5 steps if on CPU
+            if TORCH_DEVICE == "cpu" and i % 5 == 0:
+                gc.collect()
 
-    return latents
+        # Clear the alarm if set
+        if TORCH_DEVICE == "cpu":
+            signal.alarm(0)
+            
+        return latents
+        
+    except Exception as e:
+        print(f"Error in generate_image: {e}")
+        traceback.print_exc()
+        # Return empty latents as fallback
+        return torch.zeros(
+            (BATCH_SIZE, pipe.unet.config.in_channels, HEIGHT // 8, WIDTH // 8),
+            device=TORCH_DEVICE
+        )
 
-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
+def generate_images(prompt, loss_type, apply_loss, seeds, pipe, progress=gr.Progress()):
+    try:
+        images_list = []
         
-    # 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()
+        # 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
+            
+        # Process one seed at a time to save memory
+        for i, seed_no in enumerate(seeds):
+            progress((i / len(seeds)) * 0.1, desc=f"Starting seed {seed_no}")  
+            
+            # Clear memory
+            if TORCH_DEVICE == "cuda":
+                torch.cuda.empty_cache()
+            gc.collect()
+            
+            try:
+                # Generate image
+                prompts = [prompt]
+                latents = generate_image(pipe, seed_no, prompts, loss_type, loss_apply=apply_loss, progress=progress)
+                pil_images = latents_to_pil(latents, pipe)
+                images_list.extend(pil_images)
+            except Exception as e:
+                print(f"Error generating image with seed {seed_no}: {e}")
+                # Create an error image
+                error_img = Image.new('RGB', (HEIGHT, WIDTH), color=(255, 0, 0))
+                images_list.append(error_img)
+                
+            # Force garbage collection
             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]
+        # Create image grid
+        if len(images_list) > 1:
+            result = image_grid(images_list, 1, len(images_list))
+            return result
+        else:
+            return images_list[0]
+            
+    except Exception as e:
+        print(f"Error in generate_images: {e}")
+        traceback.print_exc()
+        # Create an error image
+        error_img = Image.new('RGB', (WIDTH, HEIGHT), color=(255, 0, 0))
+        return error_img
 
 # 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
@@ -231,6 +292,14 @@ def create_interface():
         Generate images using Stable Diffusion with various loss functions to guide the diffusion process.
         """)
         
+        if TORCH_DEVICE == "cpu":
+            gr.Markdown("""
+            ⚠️ **Running on CPU**: Generation will be slow and memory-intensive. 
+            Each image may take several minutes to generate.
+            """)
+        
+        pipe = None  # Initialize to None to avoid loading during interface creation
+        
         with gr.Row():
             with gr.Column():
                 prompt = gr.Textbox(
@@ -250,19 +319,44 @@ def create_interface():
                     value=False
                 )
                 
-                seeds = gr.Textbox(
-                    label="Seeds (comma-separated)",
-                    value="3000,2000,1000",
-                    lines=1
-                )
+                if TORCH_DEVICE == "cpu":
+                    seeds = gr.Textbox(
+                        label="Seeds (comma-separated) - Use fewer seeds for CPU",
+                        value="1000",
+                        lines=1
+                    )
+                else:
+                    seeds = gr.Textbox(
+                        label="Seeds (comma-separated)",
+                        value="3000,2000,1000",
+                        lines=1
+                    )
                 
-                generate_btn = gr.Button("Generate Images")
+                # Load model button
+                load_model_btn = gr.Button("Load Model")
+                model_status = gr.Textbox(label="Model Status", value="Model not loaded", interactive=False)
+                
+                generate_btn = gr.Button("Generate Images", interactive=False)
                 
             with gr.Column():
                 output_image = gr.Image(label="Generated Image")
+        
+        def load_model_fn():
+            nonlocal pipe
+            try:
+                pipe = load_model()
+                return "Model loaded successfully", True
+            except Exception as e:
+                return f"Error loading model: {str(e)}", False
                 
+        load_model_btn.click(
+            fn=load_model_fn,
+            inputs=[],
+            outputs=[model_status, generate_btn]
+        )
+        
         generate_btn.click(
-            fn=lambda p, lt, al, s: generate_images(p, lt, al, s, pipe),
+            fn=lambda p, lt, al, s, prog: generate_images(p, lt, al, s, pipe, prog),
             inputs=[prompt, loss_type, apply_loss, seeds],
             outputs=output_image
         )
@@ -275,8 +369,17 @@ def create_interface():
         - **Symmetry**: Encourages symmetrical images by minimizing differences with horizontally flipped versions
         - **Saturation**: Increases color saturation in the image
         
-        Set "N/A" and unchecks "Apply Loss Function" for normal image generation.
+        Set "N/A" and uncheck "Apply Loss Function" for normal image generation.
         """)
+        
+        if TORCH_DEVICE == "cpu":
+            gr.Markdown("""
+            ## CPU Mode Tips
+            - Use smaller prompts
+            - Process one seed at a time
+            - Be patient, generation can take 5-10 minutes per image
+            - If you encounter memory errors, try restarting the app and using even smaller dimensions
+            """)
     
     return app