fix: Refactor generator initialization and enhance mixed precision handling in process_sr function

app.py

@@ -329,7 +329,7 @@ def process_sr(
     ])
 
     seed_everything(seed)
-    generator = torch.Generator(device='cuda:1')
+    generator = torch.Generator(device=dit4sr_device)
     generator.manual_seed(seed)
 
     validation_prompt = f"{user_prompt} {positive_prompt}"
@@ -355,12 +355,19 @@ def process_sr(
     else:
         pipeline = pipeline_dit4sr_f
     
+    weight_dtype = torch.float32
+    if args.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif args.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+    
     try:
-        image = pipeline(
-            prompt=validation_prompt, control_image=input_image, num_inference_steps=num_inference_steps, generator=generator, height=height, width=width,
-            guidance_scale=cfg_scale, negative_prompt=negative_prompt, start_point=args.start_point, latent_tiled_size=args.latent_tiled_size, latent_tiled_overlap=args.latent_tiled_overlap,
-            args=args,
-        ).images[0]
+        with torch.autocast(device_type='cuda', dtype=weight_dtype, enabled=(args.mixed_precision != "no")):
+            image = pipeline(
+                prompt=validation_prompt, control_image=input_image, num_inference_steps=num_inference_steps, generator=generator, height=height, width=width,
+                guidance_scale=cfg_scale, negative_prompt=negative_prompt, start_point=args.start_point, latent_tiled_size=args.latent_tiled_size, latent_tiled_overlap=args.latent_tiled_overlap,
+                args=args,
+            ).images[0]
 
         if True:  # alpha<1.0:
             image = adain_color_fix(image, input_image)
@@ -368,7 +375,7 @@ def process_sr(
         if resize_flag:
             image = image.resize((ori_width * rscale, ori_height * rscale))
     except Exception as e:
-        print(e)
+        print(f"Error during inference: {e}")
         image = Image.new(mode="RGB", size=(512, 512))
     images.append(np.array(image))
     return images