Update app.py

app.py

@@ -16,20 +16,21 @@ pipe = DiffusionPipeline.from_pretrained(
     scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
 ).to(device)
 
-generator = th.Generator("cuda").manual_seed(0)
+generator = torch.Generator("cuda").manual_seed(0)
 
 def infer(prompt, init_image):
     init_image = Image.open(init_image).convert("RGB")
     init_image = init_image.resize((128, 128))
-    
-    res = pipe.train(
-        prompt,
-        init_image,
-        guidance_scale=7.5,
-        num_inference_steps=50,
-        generator=generator,
-        text_embedding_optimization_steps=100,
-        model_fine_tuning_optimization_steps=500)
+
+    with torch.no_grad():
+        res = pipe.train(
+            prompt,
+            init_image,
+            guidance_scale=7.5,
+            num_inference_steps=50,
+            generator=generator,
+            text_embedding_optimization_steps=100,
+            model_fine_tuning_optimization_steps=500)
     
     #with torch.no_grad():
     #    torch.cuda.empty_cache()