Update cpu fall back

app.py

@@ -26,21 +26,40 @@ pipe = None
 def load_model():
     global pipe
     if pipe is None:
-        # Use the pre-uploaded model from Hugging Face
         model_id = "Uminosachi/realisticVisionV51_v51VAE-inpainting"
-        pipe = StableDiffusionInpaintPipeline.from_pretrained(
-            model_id,
-            torch_dtype=torch.float16,
-            safety_checker=None
-        ).to("cuda")
-        pipe.enable_attention_slicing(slice_size="max")
-        pipe.enable_sequential_cpu_offload()
+        try:
+            # Try CUDA first
+            if torch.cuda.is_available():
+                device = "cuda"
+                dtype = torch.float16
+            else:
+                # Fallback to CPU
+                device = "cpu"
+                dtype = torch.float32
+            
+            pipe = StableDiffusionInpaintPipeline.from_pretrained(
+                model_id,
+                torch_dtype=dtype,
+                safety_checker=None
+            ).to(device)
+            
+            if device == "cuda":
+                pipe.enable_attention_slicing(slice_size="max")
+                pipe.enable_sequential_cpu_offload()
+            
+            print(f"Model loaded on {device}")
+            
+        except Exception as e:
+            print(f"Error loading model: {str(e)}")
+            raise
     return pipe
 
 @app.on_event("startup")
 async def startup_event():
-    if torch.cuda.is_available():
+    try:
         load_model()
+    except Exception as e:
+        print(f"Startup error: {str(e)}")
 
 def image_to_base64(image: Image.Image) -> str:
     buffered = io.BytesIO()

example_client.py

@@ -4,8 +4,8 @@ from PIL import Image
 import io
 
 def call_inpaint_api(image_path, mask_path, prompt):
-    # Update this with your actual space URL after deployment
-    url = "https://your-username-your-space-name.hf.space/inpaint"
+    # Update with your actual space URL
+    url = "https://yaseengoldfinchpc-modeltest.hf.space/inpaint"
     
     files = {
         'image': open(image_path, 'rb'),
@@ -26,4 +26,24 @@ def call_inpaint_api(image_path, mask_path, prompt):
         return 'result.png'
     else:
         print(f"Error: {response.text}")
-        return None 
+        return None
+
+def main():
+    # First test health endpoint
+    print("Testing health endpoint...")
+    response = requests.get("https://yaseengoldfinchpc-modeltest.hf.space/health")
+    print("Health Status:", response.json())
+    
+    # Test inpainting
+    print("\nTesting inpainting...")
+    # Replace these with your actual image paths
+    image_path = r"C:\Users\M. Y\Downloads\t2.png"  # Replace with your image path
+    mask_path = "generated_mask_1.png"   # Replace with your mask path
+    prompt = "add some flowers and a fountain"
+    
+    result = call_inpaint_api(image_path, mask_path, prompt)
+    if result:
+        print(f"Success! Result saved as: {result}")
+
+if __name__ == "__main__":
+    main() 

test_app.py

@@ -1,6 +1,5 @@
 from fastapi.testclient import TestClient
 from app import app
-import pytest
 
 client = TestClient(app)
 

test_local.py

@@ -0,0 +1,45 @@
+from diffusers import StableDiffusionInpaintPipeline
+import torch
+from PIL import Image
+
+def test_local():
+    # Load model
+    model_id = "Uminosachi/realisticVisionV51_v51VAE-inpainting"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float16 if device == "cuda" else torch.float32
+    
+    print(f"Using device: {device}")
+    
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(
+        model_id,
+        torch_dtype=dtype,
+        safety_checker=None
+    ).to(device)
+    
+    # Load test images
+    image_path = r"C:\Users\M. Y\Downloads\t2.png"
+    mask_path = "generated_mask_1.png"
+    
+    image = Image.open(image_path)
+    mask_image = Image.open(mask_path)
+    
+    # Resize to multiple of 8
+    width, height = (dim - dim % 8 for dim in image.size)
+    image = image.resize((width, height))
+    mask_image = mask_image.resize((width, height))
+    mask_image = mask_image.convert("L")
+    
+    # Test inference
+    result = pipe(
+        prompt="add some flowers and a fountain",
+        image=image,
+        mask_image=mask_image,
+        num_inference_steps=20,
+        guidance_scale=7.5,
+    ).images[0]
+    
+    result.save("local_test_result.png")
+    print("Test complete! Check local_test_result.png")
+
+if __name__ == "__main__":
+    test_local()