Update app.py

app.py

@@ -28,24 +28,27 @@ state_dict = load_state_dict(model_file)
 model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
     controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
 )
-# 自动检测是否有可用的 GPU，否则使用 CPU
+
+# 【核心修改 1】自动检测是否有可用的 GPU，否则使用 CPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if device == "cuda" else torch.float32
 
-# 修改你的模型加载代码
+# 修改模型加载代码，使用动态变量
 model.to(device=device, dtype=dtype)
 
+# 【核心修改 2】将 .to("cuda") 修改为动态适配变量，并根据设备调整精度
 vae = AutoencoderKL.from_pretrained(
-    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
-).to("cuda")
+    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=dtype
+).to(device)
 
+# 【核心修改 3】将 .to("cuda") 修改为动态适配变量，并根据设备调整精度与 variant
 pipe = StableDiffusionXLFillPipeline.from_pretrained(
     "SG161222/RealVisXL_V5.0_Lightning",
-    torch_dtype=torch.float16,
+    torch_dtype=dtype,
     vae=vae,
     controlnet=model,
-    variant="fp16",
-).to("cuda")
+    variant="fp16" if device == "cuda" else None,
+).to(device)
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
@@ -190,35 +193,6 @@ def infer(
     overlap_top, 
     overlap_bottom
 ):
-    """
-    Generate an outpainted image using Stable Diffusion XL with ControlNet guidance.
-    
-    This function performs intelligent image outpainting by expanding the input image 
-    according to the specified target dimensions and alignment, generating new content 
-    guided by a textual prompt. It uses a ControlNet-enabled diffusion pipeline to ensure 
-    coherent image extension.
-
-    Args:
-        image (PIL.Image): The input image to be outpainted.
-        width (int): The target width of the output image.
-        height (int): The target height of the output image.
-        overlap_percentage (int): Percentage of overlap between original and outpainted regions for seamless blending.
-        num_inference_steps (int): Number of inference steps for image generation. Higher values yield better results.
-        resize_option (str): Predefined or custom percentage to resize the input image ("Full", "50%", "33%", "25%", or "Custom").
-        custom_resize_percentage (int): Custom resize percentage if resize_option is "Custom".
-        prompt_input (str): A text prompt describing desired content for the generated region.
-        alignment (str): Alignment of the original image within the canvas ("Middle", "Left", "Right", "Top", "Bottom").
-        overlap_left (bool): Whether to allow blending on the left edge.
-        overlap_right (bool): Whether to allow blending on the right edge.
-        overlap_top (bool): Whether to allow blending on the top edge.
-        overlap_bottom (bool): Whether to allow blending on the bottom edge.
-
-    Yields:
-        Tuple[PIL.Image, PIL.Image]: 
-        - The intermediate ControlNet input image (showing the masked area).
-        - The final generated image with the inpainted region.
-    """
-    #gr.Info("10 seconds will be used from your daily ZeroGPU time credits.")
     background, mask = prepare_image_and_mask(
         image, width, height, overlap_percentage, 
         resize_option, custom_resize_percentage, alignment,
@@ -233,12 +207,13 @@ def infer(
 
     final_prompt = f"{prompt_input} , high quality, 4k"
 
+    # 【核心修改 4】将这里的编码推理设备也改用动态变量
     (
         prompt_embeds,
         negative_prompt_embeds,
         pooled_prompt_embeds,
         negative_pooled_prompt_embeds,
-    ) = pipe.encode_prompt(final_prompt, "cuda", True)
+    ) = pipe.encode_prompt(final_prompt, device, True)
 
     for image in pipe(
         prompt_embeds=prompt_embeds,
@@ -250,12 +225,6 @@ def infer(
     ):
         yield cnet_image, image
 
-    #time.sleep(1)
-    #image = image.convert("RGBA")
-    #cnet_image.paste(image, (0, 0), mask)
-
-    #return background, cnet_image
-
 
 def clear_result():
     """Clears the result ImageSlider."""
@@ -356,14 +325,14 @@ with gr.Blocks(css=css) as demo:
                                 minimum=720,
                                 maximum=1536,
                                 step=8,
-                                value=720,  # Set a default value
+                                value=720,
                             )
                             height_slider = gr.Slider(
                                 label="Target Height",
                                 minimum=720,
                                 maximum=1536,
                                 step=8,
-                                value=1280,  # Set a default value
+                                value=1280,
                             )
                         
                         num_inference_steps = gr.Slider(label="Steps", minimum=4, maximum=12, step=1, value=8)
@@ -467,47 +436,47 @@ with gr.Blocks(css=css) as demo:
         api_visibility="private"
     )
     
-    run_button.click(  # Clear the result
+    run_button.click(
         fn=clear_result,
         inputs=None,
         outputs=result,
         api_visibility="private"
-    ).then(  # Generate the new image
+    ).then(
         fn=infer,
         inputs=[input_image, width_slider, height_slider, overlap_percentage, num_inference_steps,
                 resize_option, custom_resize_percentage, prompt_input, alignment_dropdown,
                 overlap_left, overlap_right, overlap_top, overlap_bottom],
         outputs=result,
-    ).then(  # Show the "Use as Input Image" button
+    ).then(
         fn=lambda: gr.update(visible=True),
         inputs=None,
         outputs=use_as_input_button,
         api_visibility="private"
-    ).then(  # Update the history gallery
+    ).then(
         fn=lambda x, history: update_history(x[1], history),
         inputs=[result, history_gallery],
         outputs=history_gallery,
         api_visibility="private"
     )
 
-    prompt_input.submit(  # Clear the result
+    prompt_input.submit(
         fn=clear_result,
         inputs=None,
         outputs=result,
         api_visibility="private"
-    ).then(  # Generate the new image
+    ).then(
         fn=infer,
         inputs=[input_image, width_slider, height_slider, overlap_percentage, num_inference_steps,
                 resize_option, custom_resize_percentage, prompt_input, alignment_dropdown,
                 overlap_left, overlap_right, overlap_top, overlap_bottom],
         outputs=result,
         api_visibility="private"
-    ).then(  # Update the history gallery
+    ).then(
         fn=lambda x, history: update_history(x[1], history),
         inputs=[result, history_gallery],
         outputs=history_gallery,
         api_visibility="private"
-    ).then(  # Show the "Use as Input Image" button
+    ).then(
         fn=lambda: gr.update(visible=True),
         inputs=None,
         outputs=use_as_input_button,