Update app.py

app.py

@@ -38,19 +38,35 @@ def resize(value,img):
 
 def infer(source_img, prompt, guide, steps, seed, strength): 
     generator = torch.Generator('cpu').manual_seed(seed)
-    
-    source_image = resize(512, source_img)
-    source_image.save('source.png')
-    
-    images_list = img_pipe([prompt] * 1, init_image=source_image, strength=strength, guidance_scale=guide, num_inference_steps=steps)
+
+    # Ensure source_img is a valid path and load image properly
+    if isinstance(source_img, str):
+        source_image = Image.open(source_img).convert("RGB")
+    elif isinstance(source_img, Image.Image):
+        source_image = source_img.convert("RGB")
+    else:
+        raise ValueError("Unsupported input type for source_img")
+
+    # Resize image
+    source_image = source_image.resize((512, 512), Image.Resampling.LANCZOS)
+
+    # Inference
+    images_list = img_pipe(
+        [prompt],
+        init_image=source_image,
+        strength=strength,
+        guidance_scale=guide,
+        num_inference_steps=steps,
+        generator=generator
+    )
+
     images = []
-    safe_image = Image.open(r"unsafe.png")
-    
     for i, image in enumerate(images_list["images"]):
-        if(images_list["nsfw_content_detected"][i]):
-            images.append(image)
+        # if NSFW detected, show safe_image (optional)
+        if images_list["nsfw_content_detected"][i]:
+            images.append(image)  # or append `safe_image`
         else:
-            images.append(image)    
+            images.append(image)
     return images
 
 print("Great sylvain ! Everything is working fine !")