Update app.py

app.py

@@ -56,25 +56,26 @@ model = model.cuda()
 # 번역 모델 초기화
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 
+
+import torch
+from diffusers import DiffusionPipeline
+
 # FLUX 모델 설정
-device = "cpu"  # 초기에 CPU에 로드
-dtype = torch.float32  # float32 사용
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype)
+pipe = pipe.to(device)
 
-# 메모리 사용량 최적화
+# 메모리 최적화 (GPU 사용 시에만 적용)
 if torch.cuda.is_available():
-    pipe.enable_sequential_cpu_offload()
-    pipe.enable_attention_slicing(1)
+    pipe.enable_attention_slicing()
 
 @spaces.GPU(duration=300)
-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)):
+def infer_t2i(prompt, seed=42, randomize_seed=False, width=1024, height=576, 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)
-    
-    # 추론 시 GPU 사용 (가능한 경우)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    generator = torch.Generator(device=device).manual_seed(seed)
     
     with torch.no_grad():
         image = pipe(
@@ -235,7 +236,7 @@ with gr.Blocks(analytics_enabled=False, css=css) as dynamicrafter_iface:
                 t2i_randomize_seed = gr.Checkbox(label='Randomize seed', value=False)
             with gr.Row():
                 t2i_width = gr.Slider(label='Width', minimum=256, maximum=MAX_IMAGE_SIZE, step=64, value=1024)
-                t2i_height = gr.Slider(label='Height', minimum=256, maximum=MAX_IMAGE_SIZE, step=64, value=1024)
+                t2i_height = gr.Slider(label='Height', minimum=256, maximum=MAX_IMAGE_SIZE, step=64, value=576)
             with gr.Row():
                 t2i_guidance_scale = gr.Slider(label='Guidance Scale', minimum=1.0, maximum=20.0, step=0.1, value=5.0)
                 t2i_num_inference_steps = gr.Slider(label='Inference Steps', minimum=1, maximum=100, step=1, value=28)