Update app.py

app.py

@@ -54,11 +54,12 @@ model = model.cuda()
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 
 # FLUX 모델 설정
-dtype = torch.bfloat16
+dtype = torch.float16  # bfloat16 대신 float16 사용
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype)
+pipe = pipe.to(device)
 
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, device_map="auto")
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
@@ -67,17 +68,16 @@ MAX_IMAGE_SIZE = 2048
 def infer_t2i(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(seed)
+    generator = torch.Generator(device=device).manual_seed(seed)
     with torch.cuda.amp.autocast():
-        image = pipe.to("cuda")(
-            prompt = prompt, 
-            width = width,
-            height = height,
-            num_inference_steps = num_inference_steps, 
-            generator = generator,
+        image = pipe(
+            prompt=prompt, 
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps, 
+            generator=generator,
             guidance_scale=guidance_scale
         ).images[0]
-    pipe.to("cpu")
     torch.cuda.empty_cache()
     return image, seed