Update generator.py

generator.py

@@ -153,11 +153,17 @@ class RetroArtConverter:
                 if image.mode != 'RGB':
                     image = image.convert('RGB')
                 
-                # ZoeDetector handles resizing internally - just call it
-                # It returns PIL Image matching input size
+                # Get original dimensions
+                orig_width, orig_height = image.size
+                
+                # ZoeDetector handles sizing internally - just pass the image
                 depth_image = self.zoe_depth(image)
                 
-                print(f"[DEPTH] Zoe depth map generated: {image.size[0]}x{image.size[1]}")
+                # Ensure output matches original size
+                if depth_image.size != (orig_width, orig_height):
+                    depth_image = depth_image.resize((orig_width, orig_height), Image.LANCZOS)
+                
+                print(f"[DEPTH] Zoe depth map generated: {orig_width}x{orig_height}")
                 return depth_image
                 
             except Exception as e:
@@ -595,57 +601,22 @@ class RetroArtConverter:
             try:
                 print("Encoding prompts with Compel...")
                 
-                # Encode prompts
+                # Encode prompts with Compel
                 conditioning = self.compel(prompt)
                 negative_conditioning = self.compel(negative_prompt)
                 
-                # Extract embeddings - Compel returns (prompt_embeds, pooled_embeds)
-                prompt_embeds = conditioning[0]
-                pooled_prompt_embeds = conditioning[1]
-                negative_prompt_embeds = negative_conditioning[0]
-                negative_pooled_prompt_embeds = negative_conditioning[1]
-                
-                # Ensure consistent shapes (SDXL uses 77 tokens max)
-                max_length = max(prompt_embeds.shape[1], negative_prompt_embeds.shape[1])
-                
-                # Pad if needed
-                if prompt_embeds.shape[1] < max_length:
-                    padding = torch.zeros(
-                        prompt_embeds.shape[0],
-                        max_length - prompt_embeds.shape[1],
-                        prompt_embeds.shape[2],
-                        dtype=prompt_embeds.dtype,
-                        device=prompt_embeds.device
-                    )
-                    prompt_embeds = torch.cat([prompt_embeds, padding], dim=1)
-                
-                if negative_prompt_embeds.shape[1] < max_length:
-                    padding = torch.zeros(
-                        negative_prompt_embeds.shape[0],
-                        max_length - negative_prompt_embeds.shape[1],
-                        negative_prompt_embeds.shape[2],
-                        dtype=negative_prompt_embeds.dtype,
-                        device=negative_prompt_embeds.device
-                    )
-                    negative_prompt_embeds = torch.cat([negative_prompt_embeds, padding], dim=1)
-                
-                # Truncate if needed
-                if prompt_embeds.shape[1] > 77:
-                    prompt_embeds = prompt_embeds[:, :77, :]
-                if negative_prompt_embeds.shape[1] > 77:
-                    negative_prompt_embeds = negative_prompt_embeds[:, :77, :]
-                
-                pipe_kwargs["prompt_embeds"] = prompt_embeds
-                pipe_kwargs["pooled_prompt_embeds"] = pooled_prompt_embeds
-                pipe_kwargs["negative_prompt_embeds"] = negative_prompt_embeds
-                pipe_kwargs["negative_pooled_prompt_embeds"] = negative_pooled_prompt_embeds
+                # Extract embeddings
+                pipe_kwargs["prompt_embeds"] = conditioning[0]
+                pipe_kwargs["pooled_prompt_embeds"] = conditioning[1]
+                pipe_kwargs["negative_prompt_embeds"] = negative_conditioning[0]
+                pipe_kwargs["negative_pooled_prompt_embeds"] = negative_conditioning[1]
                 
                 compel_success = True
-                print(f"[OK] Compel encoded: pos={prompt_embeds.shape}, neg={negative_prompt_embeds.shape}")
+                print("[OK] Using Compel-encoded prompts")
                 
             except Exception as e:
-                print(f"[COMPEL] Failed: {e}")
-                print("[COMPEL] Falling back to standard encoding")
+                print(f"[COMPEL] Encoding failed: {e}")
+                print(f"[COMPEL] Using standard prompt encoding instead")
                 compel_success = False
         
         # Use standard prompts if Compel failed or not available