Update app.py

app.py

@@ -22,11 +22,11 @@ from diffusers import FluxPipeline
 # 상단에 import 추가
 from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 
-# 번역 모델 초기화
+
 model_name = "Helsinki-NLP/opus-mt-ko-en"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-translator = pipeline("translation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to('cpu')  # CPU로 강제 지정
+translator = pipeline("translation", model=model, tokenizer=tokenizer, device=-1)  # CPU 사용
 
 
 def translate_to_english(text: str) -> str:
@@ -168,18 +168,46 @@ def generate_background(prompt: str, aspect_ratio: str) -> Image.Image:
         # 8의 배수로 조정
         width, height = adjust_size_to_multiple_of_8(width, height)
         
+        # 프롬프트 전처리
+        if not prompt or prompt.strip() == "":
+            prompt = "plain white background"
+        
         with timer("Background generation"):
-            image = pipe(
-                prompt=prompt,
-                width=width,
-                height=height,
-                num_inference_steps=8,
-                guidance_scale=4.0,
-            ).images[0]
+            try:
+                image = pipe(
+                    prompt=prompt,
+                    width=width,
+                    height=height,
+                    num_inference_steps=8,
+                    guidance_scale=4.0,
+                    max_length=77,  # CLIP 텍스트 인코더의 최대 길이 제한
+                ).images[0]
+            except Exception as e:
+                print(f"Pipeline error: {str(e)}")
+                # 오류 발생 시 기본 흰색 배경 생성
+                image = Image.new('RGB', (width, height), 'white')
                 
         return image
     except Exception as e:
-        raise gr.Error(f"Background generation failed: {str(e)}")
+        print(f"Background generation error: {str(e)}")
+        # 최후의 폴백: 기본 흰색 배경 반환
+        return Image.new('RGB', (512, 512), 'white')
+
+# FLUX 파이프라인 초기화 부분 수정
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    torch_dtype=torch.float32,  # bfloat16 대신 float32 사용
+    use_auth_token=HF_TOKEN
+)
+pipe.load_lora_weights(
+    hf_hub_download(
+        "ByteDance/Hyper-SD",
+        "Hyper-FLUX.1-dev-8steps-lora.safetensors",
+        use_auth_token=HF_TOKEN
+    )
+)
+pipe.fuse_lora(lora_scale=0.125)
+pipe.to(device=device)
 
 
 def create_position_grid():
@@ -302,30 +330,36 @@ def process_prompt(img: Image.Image, prompt: str, bg_prompt: str | None = None,
         if img is None or prompt.strip() == "":
             raise gr.Error("Please provide both image and prompt")
         
-        print(f"Processing with position: {position}, scale: {scale_percent}")  # 디버그 로그
+        print(f"Processing with position: {position}, scale: {scale_percent}")
         
-        # 프롬프트 번역
-        prompt = translate_to_english(prompt)
-        if bg_prompt:
-            bg_prompt = translate_to_english(bg_prompt)
+        try:
+            # 프롬프트 번역 시도
+            prompt = translate_to_english(prompt)
+            if bg_prompt:
+                bg_prompt = translate_to_english(bg_prompt)
+        except Exception as e:
+            print(f"Translation error (continuing with original text): {str(e)}")
         
         # Process the image
         results, _ = _process(img, prompt, bg_prompt, aspect_ratio)
         
         if bg_prompt:
-            # 여기서 한 번만 합성 수행
-            combined = combine_with_background(
-                foreground=results[2],  # 추출된 오브젝트
-                background=results[1],  # 생성된 배경
-                position=position,
-                scale_percent=scale_percent
-            )
-            print(f"Combined image created with position: {position}")  # 디버그 로그
-            return combined, results[2]
+            try:
+                combined = combine_with_background(
+                    foreground=results[2],
+                    background=results[1],
+                    position=position,
+                    scale_percent=scale_percent
+                )
+                print(f"Combined image created with position: {position}")
+                return combined, results[2]
+            except Exception as e:
+                print(f"Combination error: {str(e)}")
+                return results[1], results[2]
         
         return results[1], results[2]
     except Exception as e:
-        print(f"Error in process_prompt: {str(e)}")  # 디버그 로그
+        print(f"Error in process_prompt: {str(e)}")
         raise gr.Error(str(e))
 
 def process_bbox(img: Image.Image, box_input: str) -> tuple[Image.Image, Image.Image]: