Update app.py

app.py

@@ -39,81 +39,48 @@ style_token_dict = {
     "Birb Style": '<birb-style>'
 }
 
-def apply_guidance(latents, guidance_method, loss_scale):
+def apply_guidance(image, guidance_method, loss_scale):
+    # Convert PIL Image to tensor
+    img_tensor = tfms.ToTensor()(image).unsqueeze(0).to(torch_device)
+    
     if guidance_method == 'Grayscale':
-        rgb = latents_to_pil(latents)[0]
-        gray = rgb.convert('L')
-        gray_latents = pil_to_latent(gray.convert('RGB'))
-        return latents + (gray_latents - latents) * loss_scale
+        gray = tfms.Grayscale(3)(img_tensor)
+        guided = img_tensor + (gray - img_tensor) * (loss_scale / 10000)
     elif guidance_method == 'Bright':
-        bright_latents = F.relu(latents)  # Simple brightness increase
-        return latents + (bright_latents - latents) * loss_scale
+        bright = F.relu(img_tensor)  # Simple brightness increase
+        guided = img_tensor + (bright - img_tensor) * (loss_scale / 10000)
     elif guidance_method == 'Contrast':
-        mean = latents.mean()
-        contrast_latents = (latents - mean) * 2 + mean
-        return latents + (contrast_latents - latents) * loss_scale
+        mean = img_tensor.mean()
+        contrast = (img_tensor - mean) * 2 + mean
+        guided = img_tensor + (contrast - img_tensor) * (loss_scale / 10000)
     elif guidance_method == 'Symmetry':
-        flipped_latents = torch.flip(latents, [3])  # Flip horizontally
-        return latents + (flipped_latents - latents) * loss_scale
+        flipped = torch.flip(img_tensor, [3])  # Flip horizontally
+        guided = img_tensor + (flipped - img_tensor) * (loss_scale / 10000)
     elif guidance_method == 'Saturation':
-        rgb = latents_to_pil(latents)[0]
-        saturated = tfms.functional.adjust_saturation(tfms.ToTensor()(rgb), 2)
-        saturated_latents = pil_to_latent(tfms.ToPILImage()(saturated))
-        return latents + (saturated_latents - latents) * loss_scale
+        saturated = tfms.functional.adjust_saturation(img_tensor, 2)
+        guided = img_tensor + (saturated - img_tensor) * (loss_scale / 10000)
     else:
-        return latents
+        return image
+
+    # Convert back to PIL Image
+    guided = guided.squeeze(0).clamp(0, 1)
+    guided = (guided * 255).byte().cpu().permute(1, 2, 0).numpy()
+    return Image.fromarray(guided)
 
 def generate_with_guidance(prompt, num_inference_steps, guidance_scale, seed, guidance_method, loss_scale):
+    # Generate image with pipeline
     generator = torch.Generator(device=torch_device).manual_seed(seed)
+    image = sd_pipeline(
+        prompt,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        generator=generator
+    ).images[0]
+
+    # Apply guidance
+    guided_image = apply_guidance(image, guidance_method, loss_scale)
     
-    # Get the text embeddings
-    text_input = sd_pipeline.tokenizer(prompt, padding="max_length", max_length=sd_pipeline.tokenizer.model_max_length, truncation=True, return_tensors="pt")
-    with torch.no_grad():
-        text_embeddings = sd_pipeline.text_encoder(text_input.input_ids.to(torch_device))[0]
-    
-    # Set the timesteps
-    sd_pipeline.scheduler.set_timesteps(num_inference_steps)
-    
-    # Prepare latents
-    latents = torch.randn(
-        (1, sd_pipeline.unet.in_channels, 64, 64),
-        generator=generator,
-        device=torch_device
-    )
-    latents = latents * sd_pipeline.scheduler.init_noise_sigma
-    
-    # Denoising loop
-    for t in tqdm(sd_pipeline.scheduler.timesteps):
-        # Expand the latents for classifier-free guidance
-        latent_model_input = torch.cat([latents] * 2)
-        latent_model_input = sd_pipeline.scheduler.scale_model_input(latent_model_input, timestep=t)
-        
-        # Predict the noise residual
-        with torch.no_grad():
-            noise_pred = sd_pipeline.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
-        
-        # Perform guidance
-        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-        
-        # Apply custom guidance
-        latents = apply_guidance(latents, guidance_method, loss_scale / 10000)  # Normalize loss_scale
-        
-        # Compute the previous noisy sample x_t -> x_t-1
-        latents = sd_pipeline.scheduler.step(noise_pred, t, latents).prev_sample
-    
-    # Scale and decode the image latents with vae
-    latents = 1 / 0.18215 * latents
-    with torch.no_grad():
-        image = sd_pipeline.vae.decode(latents).sample
-    
-    # Convert to PIL Image
-    image = (image / 2 + 0.5).clamp(0, 1)
-    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
-    image = (image * 255).round().astype("uint8")[0]
-    image = Image.fromarray(image)
-    
-    return image
+    return guided_image
 
 def inference(text, style, inference_step, guidance_scale, seed, guidance_method, loss_scale):
     prompt = text + " " + style_token_dict[style]