Update app.py

app.py

@@ -48,28 +48,75 @@ import cv2
 import numpy as np
 import sys
 import io
+
 # FluxPipeline import 부분을 수정
 from diffusers import StableDiffusionPipeline, DiffusionPipeline
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers import AutoPipelineForText2Image
+
+# Model initialization 부분 수정
+if not path.exists(cache_path):
+    os.makedirs(cache_path, exist_ok=True)
 
-# 기존의 pipe 초기화 부분을 다음과 같이 수정
-pipe = StableDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-2-1",
+# 파이프라인 초기화 수정
+pipe = AutoPipelineForText2Image.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
     torch_dtype=torch.float16,
-    safety_checker=None
+    use_safetensors=True,
+    variant="fp16"
 )
 pipe.to("cuda")
 
-# LoRA 가중치 로드
-pipe.load_lora_weights(
-    hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors")
-)
-pipe.fuse_lora(lora_scale=0.125)
-
 # 안전 검사기 설정
 pipe.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretrained(
     "CompVis/stable-diffusion-safety-checker"
 )
+
+# process_and_save_image 함수 수정
+@spaces.GPU
+def process_and_save_image(height, width, steps, scales, prompt, seed):
+    is_safe, processed_prompt = process_prompt(prompt)
+    if not is_safe:
+        gr.Warning("부적절한 내용이 포함된 프롬프트입니다.")
+        return None, load_gallery()
+          
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
+        try:
+            generated_image = pipe(
+                prompt=processed_prompt,
+                negative_prompt="low quality, worst quality, bad anatomy, bad composition, poor, low effort",
+                num_inference_steps=steps,
+                guidance_scale=scales,
+                height=height,
+                width=width,
+                generator=torch.Generator("cuda").manual_seed(int(seed))
+            ).images[0]
+                    
+            # PIL Image로 확실하게 변환
+            if not isinstance(generated_image, Image.Image):
+                generated_image = Image.fromarray(generated_image)
+                    
+            # RGB 모드로 변환
+            if generated_image.mode != 'RGB':
+                generated_image = generated_image.convert('RGB')
+                    
+            # 메모리에서 PNG로 변환
+            img_byte_arr = io.BytesIO()
+            generated_image.save(img_byte_arr, format='PNG')
+            img_byte_arr = img_byte_arr.getvalue()
+                
+            # 디스크에 저장
+            saved_path = save_image(generated_image)
+            if saved_path is None:
+                logger.warning("Failed to save generated image")
+                return None, load_gallery()
+                    
+            # PNG 형식으로 다시 로드
+            return Image.open(io.BytesIO(img_byte_arr)), load_gallery()
+        except Exception as e:
+            logger.error(f"Error in image generation: {str(e)}")
+            return None, load_gallery()
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)