Update app.py

app.py

@@ -43,16 +43,31 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
-# 텍스트 인코더를 float16으로 강제 변환
-pipe.text_encoder = pipe.text_encoder.to("cuda", dtype=torch.float16)
-# 만약 text_projection 속성이 있다면, 입력이 float16이 아니면 half로 캐스팅하도록 오버라이딩
-if hasattr(pipe.text_encoder, "text_projection"):
-    original_text_projection_forward = pipe.text_encoder.text_projection.forward
-    def fixed_text_projection_forward(x):
-        if x.dtype != torch.float16:
-            x = x.half()
-        return original_text_projection_forward(x)
-    pipe.text_encoder.text_projection.forward = fixed_text_projection_forward
+# 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)
 
 def can_expand(source_width, source_height, target_width, target_height, alignment):
     """Checks if the image can be expanded based on the alignment."""
@@ -153,24 +168,25 @@ def infer(image, width, height, overlap_width, num_inference_steps, resize_optio
     cnet_image = background.copy()
     cnet_image.paste(0, (0, 0), mask)
 
-    final_prompt = f"{prompt_input} , high quality, 4k"
-
-    (
-        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
+    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
 
     image = image.convert("RGBA")
     cnet_image.paste(image, (0, 0), mask)
@@ -371,4 +387,4 @@ with gr.Blocks(css=css) as demo:
         outputs=use_as_input_button,
     )
 
-demo.queue(max_size=12).launch(share=False)
+demo.queue(max_size=12).launch(share=False)