Update app.py

app.py

@@ -1,9 +1,6 @@
-#!/usr/bin/env python
-
 import os
 import random
 import uuid
-
 import gradio as gr
 import numpy as np
 from PIL import Image
@@ -13,18 +10,20 @@ from diffusers import DiffusionPipeline
 
 DESCRIPTION = """# Playground v2.5"""
 if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"
+    DESCRIPTION += "\n<p>CPU에서 실행 중 🥶 이 데모는 CPU에서 작동하지 않을 수 있습니다.</p>"
 
+# 일반 설정
 MAX_SEED = np.iinfo(np.int32).max
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1536"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 
+# 디바이스 설정
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-NUM_IMAGES_PER_PROMPT = 1
-
+# CUDA가 사용 가능할 경우 글로벌로 파이프라인 로드
+pipe = None
 if torch.cuda.is_available():
     pipe = DiffusionPipeline.from_pretrained(
         "playgroundai/playground-v2.5-1024px-aesthetic",
@@ -32,49 +31,34 @@ if torch.cuda.is_available():
         use_safetensors=True,
         add_watermarker=False,
         variant="fp16"
-    )
+    ).to(device)
+
     if ENABLE_CPU_OFFLOAD:
         pipe.enable_model_cpu_offload()
-    else:
-        pipe.to(device)    
-        print("Loaded on Device!")
     
     if USE_TORCH_COMPILE:
         pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-        print("Model Compiled!")
-
 
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
     return unique_name
 
-
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-
 @spaces.GPU(enable_queue=True)
-def generate(
-    prompt: str,
-    negative_prompt: str = "",
-    use_negative_prompt: bool = False,
-    seed: int = 0,
-    width: int = 1024,
-    height: int = 1024,
-    guidance_scale: float = 3,
-    randomize_seed: bool = False,
-    use_resolution_binning: bool = True,
-    progress=gr.Progress(track_tqdm=True),
-):
-    global pipe  # 전역 변수로서의 pipe를 참조하도록 합니다.
-    
-    if not torch.cuda.is_available():
-        return [], "GPU가 사용 불가능합니다. 이 데모는 GPU에서만 실행될 수 있습니다."
-
-    # GPU 사용 가능한 경우의 로직을 계속 진행합니다.
+def generate(prompt: str, negative_prompt: str = "", use_negative_prompt: bool = False,
+             seed: int = 0, width: int = 1024, height: int = 1024, guidance_scale: float = 7.0,
+             randomize_seed: bool = False, use_resolution_binning: bool = True,
+             progress=gr.Progress(track_tqdm=True)):
+    global pipe
+    if pipe is None:
+        raise Exception("모델 파이프라인이 로드되지 않았습니다. 이 데모는 GPU가 필요합니다.")
+
+    # pipe가 올바른 디바이스에 있는지 확인
     pipe.to(device)
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator().manual_seed(seed)