Update app.py

app.py

@@ -70,22 +70,21 @@ try:
     pipe = FluxPipeline.from_pretrained(
         "black-forest-labs/FLUX.1-dev",
         torch_dtype=torch.float16,
-        use_auth_token=HF_TOKEN,
-        safety_checker=None,
-        device_map="balanced"  # 'auto' 대신 'balanced' 사용
+        use_auth_token=HF_TOKEN
     )
     print("FLUX pipeline initialized successfully")
     
     # 메모리 최적화 설정
     pipe.enable_attention_slicing(slice_size=1)
-    pipe.enable_model_cpu_offload()
-    print("Pipeline optimization settings applied")
-
-    # 추가 메모리 최적화
+    
+    # GPU 설정
     if torch.cuda.is_available():
+        pipe = pipe.to("cuda:0")
         torch.cuda.empty_cache()
         torch.backends.cudnn.benchmark = True
         torch.backends.cuda.matmul.allow_tf32 = True
+    
+    print("Pipeline optimization settings applied")
 
 except Exception as e:
     print(f"Error initializing FLUX pipeline: {str(e)}")