Update app.py

app.py

@@ -81,10 +81,14 @@ def initialize_model():
             lora_path = hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors")
             pipe.load_lora_weights(lora_path)
             pipe.fuse_lora(lora_scale=0.125)
+            
+            # 주의: 여기서 device를 명시적으로 지정 (모든 컴포넌트에 적용)
             pipe.to(device="cuda", dtype=torch.bfloat16)
             
-            # 안전 검사기 추가
+            # 안전 검사기 추가 및 올바른 장치로 이동
             pipe.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
+            if hasattr(pipe, 'safety_checker') and pipe.safety_checker is not None:
+                pipe.safety_checker.to("cuda")
             
         print("모델 로딩 완료")
         return True
@@ -270,14 +274,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     def process_image(height, width, steps, scales, prompt, seed):
         global pipe
         
-        # 모델 초기화 상태 확인
-        if pipe is None:
-            return None, "모델을 로딩 중입니다... 처음 실행 시 시간이 소요될 수 있습니다.", True, "", False
-            
-            model_loaded = initialize_model()
-            if not model_loaded:
-                return None, "", False, "모델 로딩 중 오류가 발생했습니다. 페이지를 새로고침하고 다시 시도해 주세요.", True
-        
         # 입력값 검증
         if not prompt or prompt.strip() == "":
             return None, "", False, "이미지 설명을 입력해주세요.", True
@@ -298,21 +294,25 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             else:
                 seed = int(seed)  # 타입 변환 안전하게 처리
                 
-            # 이미지 생성 상태 메시지
-            loading_message = "이미지를 생성 중입니다..."
-                
+            # 높이와 너비를 64의 배수로 조정 (FLUX 모델 요구사항)
+            height = (int(height) // 64) * 64
+            width = (int(width) // 64) * 64
+            
+            # 안전장치 - 최대값 제한
+            steps = min(int(steps), 25)
+            scales = max(min(float(scales), 5.0), 0.0)
+            
             # 이미지 생성
             with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
+                # 중요: generator 설정 시 device를 명시적으로 지정
                 generator = torch.Generator(device="cuda").manual_seed(seed)
                 
-                # 높이와 너비를 64의 배수로 조정 (FLUX 모델 요구사항)
-                height = (int(height) // 64) * 64
-                width = (int(width) // 64) * 64
-                
-                # 안전장치 - 최대값 제한
-                steps = min(int(steps), 25)
-                scales = max(min(float(scales), 5.0), 0.0)
-                
+                # 모든 텐서가 같은 디바이스에 있는지 확인
+                for name, module in pipe.components.items():
+                    if hasattr(module, 'device') and module.device.type != "cuda":
+                        module.to("cuda")
+                        
+                # 이미지 생성 - 모든 매개변수에 device를 명시적 지정
                 generated_image = pipe(
                     prompt=[filtered_prompt],
                     generator=generator,
@@ -320,10 +320,10 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     guidance_scale=scales,
                     height=height,
                     width=width,
-                    max_sequence_length=256
+                    max_sequence_length=256,
+                    device="cuda"  # 명시적 device 지정
                 ).images[0]
                 
-                # 성공 시 이미지 반환, 상태 메시지 숨김
                 return generated_image, "", False, "", False
                 
         except Exception as e:
@@ -342,11 +342,19 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     
     # 이미지 생성 준비 함수
     def prepare_generation(height, width, steps, scales, prompt, seed):
+        global pipe
+        
         # 모델이 아직 로드되지 않았다면 로드
         if pipe is None:
+            # 로딩 상태 표시
+            loading_message = "모델을 로딩 중입니다... 처음 실행 시 시간이 소요될 수 있습니다."
+            
             is_loaded = initialize_model()
             if not is_loaded:
-                return None, "모델 로딩에 실패했습니다. 페이지를 새로고침하고 다시 시도해 주세요.", True, "", False
+                return None, "", False, "모델 로딩에 실패했습니다. 페이지를 새로고침하고 다시 시도해 주세요.", True
+            
+        # 생성 상태 표시
+        loading_message = "이미지를 생성 중입니다..."
             
         # 생성 프로세스 시작
         return process_image(height, width, steps, scales, prompt, seed)
@@ -364,5 +372,5 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     )
 
 if __name__ == "__main__":
-    # 앱 시작 시 모델 미리 로드하지 않음 (첫 요청 시 지연 로딩)
+
     demo.queue(max_size=10).launch()