Update app.py

app.py

@@ -12,6 +12,14 @@ from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 from PIL import Image, ImageDraw
 import numpy as np
 
+# Initialize CUDA if available
+if torch.cuda.is_available():
+    device = "cuda"
+    dtype = torch.float16
+else:
+    device = "cpu"
+    dtype = torch.float32
+
 config_file = hf_hub_download(
     "xinsir/controlnet-union-sdxl-1.0",
     filename="config_promax.json",
@@ -37,20 +45,19 @@ result = ControlNetModel_Union._load_pretrained_model(
 
 # Use the first element from the result
 model = result[0]
-model = model.to(device="cuda", dtype=torch.float16)
-
+model = model.to(device=device, dtype=dtype)
 
 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)
 
 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 dtype == torch.float16 else None,
+).to(device)
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
@@ -152,7 +159,6 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
     elif alignment == "Bottom":
         bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height
 
-
     # Draw the mask
     mask_draw.rectangle([
         (left_overlap, top_overlap),
@@ -181,39 +187,45 @@ def preview_image_and_mask(image, width, height, overlap_percentage, resize_opti
 
 @spaces.GPU(duration=24)
 def infer(image, width, height, overlap_percentage, num_inference_steps, resize_option, custom_resize_percentage, prompt_input, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
-    background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)
-    
-    if not can_expand(background.width, background.height, width, height, alignment):
-        alignment = "Middle"
-
-    cnet_image = background.copy()
-    cnet_image.paste(0, (0, 0), mask)
-
-    final_prompt = f"{prompt_input} , high quality, 4k" if prompt_input else "high quality, 4k"
-
-    # Use with torch.autocast to ensure consistent dtype
-    with torch.autocast(device_type="cuda", 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)
-
-    yield background, cnet_image
+    try:
+        background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)
+        
+        if not can_expand(background.width, background.height, width, height, alignment):
+            alignment = "Middle"
+
+        cnet_image = background.copy()
+        cnet_image.paste(0, (0, 0), mask)
+
+        final_prompt = f"{prompt_input} , high quality, 4k" if prompt_input else "high quality, 4k"
+
+        # Use with torch.autocast to ensure consistent dtype
+        with torch.autocast(device_type=device, dtype=dtype):
+            (
+                prompt_embeds,
+                negative_prompt_embeds,
+                pooled_prompt_embeds,
+                negative_pooled_prompt_embeds,
+            ) = pipe.encode_prompt(final_prompt, device, 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)
+
+        yield background, cnet_image
+        
+    except Exception as e:
+        print(f"Error in infer function: {e}")
+        # Return a placeholder or error image
+        yield None, None
 
 def clear_result():
     """Clears the result ImageSlider."""
@@ -253,9 +265,21 @@ def update_history(new_image, history):
     """Updates the history gallery with the new image."""
     if history is None:
         history = []
-    history.insert(0, new_image)
+    
+    # Check if new_image is valid and has the expected structure
+    if new_image is not None and isinstance(new_image, (tuple, list)) and len(new_image) > 1:
+        if new_image[1] is not None:  # Check if the second element exists
+            history.insert(0, new_image[1])
+    
     return history
 
+# Safe wrapper for update_history to handle None values
+def safe_update_history(result, history):
+    """Safely updates the history gallery with the new image."""
+    if result is None:
+        return history
+    return update_history(result, history)
+
 css = """
 .gradio-container {
     width: 1200px !important;
@@ -358,8 +382,6 @@ with gr.Blocks(theme="soft", css=css) as demo:
                     inputs=[input_image, width_slider, height_slider, alignment_dropdown],
                 )
 
-                
-
             with gr.Column():
                 result = ImageSlider(
                     interactive=False,
@@ -370,11 +392,11 @@ with gr.Blocks(theme="soft", css=css) as demo:
                 history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
                 preview_image = gr.Image(label="Preview")
 
-        
-
     def use_output_as_input(output_image):
         """Sets the generated output as the new input image."""
-        return gr.update(value=output_image[1])
+        if output_image is not None and isinstance(output_image, (tuple, list)) and len(output_image) > 1:
+            return gr.update(value=output_image[1])
+        return gr.update()
 
     use_as_input_button.click(
         fn=use_output_as_input,
@@ -421,7 +443,7 @@ with gr.Blocks(theme="soft", css=css) as demo:
                 overlap_left, overlap_right, overlap_top, overlap_bottom],
         outputs=result,
     ).then(  # Update the history gallery
-        fn=lambda x, history: update_history(x[1], history),
+        fn=safe_update_history,
         inputs=[result, history_gallery],
         outputs=history_gallery,
     ).then(  # Show the "Use as Input Image" button
@@ -441,7 +463,7 @@ with gr.Blocks(theme="soft", css=css) as demo:
                 overlap_left, overlap_right, overlap_top, overlap_bottom],
         outputs=result,
     ).then(  # Update the history gallery
-        fn=lambda x, history: update_history(x[1], history),
+        fn=safe_update_history,
         inputs=[result, history_gallery],
         outputs=history_gallery,
     ).then(  # Show the "Use as Input Image" button