Update app.py

app.py

@@ -5,12 +5,13 @@ import torch
 import os
 from diffusers import DiffusionPipeline, DDIMScheduler
 MY_SECRET_TOKEN=os.environ.get('HF_TOKEN_SD')
+
 has_cuda = torch.cuda.is_available()
-device = 'cpu' if not has_cuda else 'cuda'
+device = torch.device('cpu' if not has_cuda else 'cuda')
 
 pipe = DiffusionPipeline.from_pretrained(
     "CompVis/stable-diffusion-v1-4",
-        safety_checker=None,
+    safety_checker=None,
     use_auth_token=MY_SECRET_TOKEN,
     custom_pipeline="imagic_stable_diffusion",
     scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
@@ -19,6 +20,9 @@ pipe = DiffusionPipeline.from_pretrained(
 #generator = torch.Generator("cuda").manual_seed(0)
 
 def infer(prompt, init_image):
+    init_image = Image.open(BytesIO(init_image)).convert("RGB")
+    init_image = init_image.resize((512, 512))
+    
     res = pipe.train(
         prompt,
         init_image,