Update app.py

app.py

@@ -17,13 +17,15 @@ import pandas as pd
 from transformers import pipeline
 
 import logging
-import random
 import warnings
 import numpy as np
 from diffusers import FluxControlNetModel
 from diffusers.pipelines import FluxControlNetPipeline
 from PIL import Image
 from huggingface_hub import snapshot_download
+from gradio_imageslider import ImageSlider
+
+
 
 # 번역 모델 로드
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
@@ -488,9 +490,70 @@ css = '''
 #component-11{align-self: stretch;}
 footer {visibility: hidden;}
 '''
-# 업스케일 함수 추가
+
+
+
+huggingface_token = os.getenv("HUGGINFACE_TOKEN")
+
+model_path = snapshot_download(
+    repo_id="black-forest-labs/FLUX.1-dev", 
+    repo_type="model", 
+    ignore_patterns=["*.md", "*..gitattributes"],
+    local_dir="FLUX.1-dev",
+    token=huggingface_token, # type a new token-id.
+)
+
+
+# Load pipeline
+controlnet = FluxControlNetModel.from_pretrained(
+    "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
+).to(device)
+pipe = FluxControlNetPipeline.from_pretrained(
+    model_path, controlnet=controlnet, torch_dtype=torch.bfloat16
+)
+pipe.to(device)
+
+MAX_SEED = 1000000
+MAX_PIXEL_BUDGET = 1024 * 1024
+
+
+
+def process_input(input_image, upscale_factor):
+    w, h = input_image.size
+    w_original, h_original = w, h
+    aspect_ratio = w / h
+
+    was_resized = False
+
+    if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
+        warnings.warn(
+            f"Requested output image is too large ({w * upscale_factor}x{h * upscale_factor}). Resizing to ({int(aspect_ratio * MAX_PIXEL_BUDGET ** 0.5 // upscale_factor), int(MAX_PIXEL_BUDGET ** 0.5 // aspect_ratio // upscale_factor)}) pixels."
+        )
+        gr.Info(
+            f"Requested output image is too large ({w * upscale_factor}x{h * upscale_factor}). Resizing input to ({int(aspect_ratio * MAX_PIXEL_BUDGET ** 0.5 // upscale_factor), int(MAX_PIXEL_BUDGET ** 0.5 // aspect_ratio // upscale_factor)}) pixels budget."
+        )
+        input_image = input_image.resize(
+            (
+                int(aspect_ratio * MAX_PIXEL_BUDGET**0.5 // upscale_factor),
+                int(MAX_PIXEL_BUDGET**0.5 // aspect_ratio // upscale_factor),
+            )
+        )
+        was_resized = True
+
+    # resize to multiple of 8
+    w, h = input_image.size
+    w = w - w % 8
+    h = h - h % 8
+
+    return input_image.resize((w, h)), w_original, h_original, was_resized
+
+MAX_PIXEL_BUDGET = 1024 * 1024
+    
 @spaces.GPU
 def upscale(input_image, progress=gr.Progress(track_tqdm=True)):
+    if input_image is None:
+        raise gr.Error("No image to upscale. Please generate an image first.")
+    
     # 입력 이미지 처리
     input_image, w_original, h_original, was_resized = process_input(input_image, 4)
     
@@ -500,9 +563,9 @@ def upscale(input_image, progress=gr.Progress(track_tqdm=True)):
     generator = torch.Generator(device=device).manual_seed(random.randint(0, MAX_SEED))
 
     gr.Info("Upscaling image to 4096x4096...")
-    image = pipe(
+    image = pipe_controlnet(
         prompt="",
-        control_image=control_image,
+        image=control_image,
         controlnet_conditioning_scale=0.6,
         num_inference_steps=28,
         guidance_scale=3.5,