Update app.py

app.py

@@ -43,31 +43,8 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
-# Make sure all text encoder components use the same dtype
-pipe.text_encoder = pipe.text_encoder.to(dtype=torch.float16)
-pipe.text_encoder_2 = pipe.text_encoder_2.to(dtype=torch.float16)
-
-# Patch the text encoder forward methods to ensure consistent dtype
-def patch_text_encoder_forward(encoder):
-    original_forward = encoder.forward
-    
-    def patched_forward(*args, **kwargs):
-        # Convert input tensors to float16
-        if len(args) > 0 and isinstance(args[0], torch.Tensor):
-            args = list(args)
-            args[0] = args[0].to(dtype=torch.float16)
-        
-        for key in kwargs:
-            if isinstance(kwargs[key], torch.Tensor):
-                kwargs[key] = kwargs[key].to(dtype=torch.float16)
-        
-        return original_forward(*args, **kwargs)
-    
-    encoder.forward = patched_forward
-
-# Apply the patch to both encoders
-patch_text_encoder_forward(pipe.text_encoder)
-patch_text_encoder_forward(pipe.text_encoder_2)
+# The key differences are below - don't modify the text encoder directly
+# We'll fix it in the pipeline code instead
 
 def can_expand(source_width, source_height, target_width, target_height, alignment):
     """Checks if the image can be expanded based on the alignment."""
@@ -170,23 +147,19 @@ def infer(image, width, height, overlap_width, num_inference_steps, resize_optio
 
     final_prompt = f"{prompt_input} , high quality, 4k" if prompt_input else "high quality, 4k"
 
-    with torch.cuda.amp.autocast(dtype=torch.float16):
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        ) = pipe.encode_prompt(final_prompt, "cuda", True)
-
-        for image in pipe(
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
-            image=cnet_image,
-            num_inference_steps=num_inference_steps
-        ):
-            yield cnet_image, image
+    # Important fix: Use the original pipeline's method without modifications
+    # Let it handle the dtype conversions internally
+    encoded_prompts = pipe.encode_prompt(final_prompt, "cuda", True)
+    
+    for image in pipe(
+        prompt_embeds=encoded_prompts[0],
+        negative_prompt_embeds=encoded_prompts[1],
+        pooled_prompt_embeds=encoded_prompts[2],
+        negative_pooled_prompt_embeds=encoded_prompts[3],
+        image=cnet_image,
+        num_inference_steps=num_inference_steps
+    ):
+        yield cnet_image, image
 
     image = image.convert("RGBA")
     cnet_image.paste(image, (0, 0), mask)