Update app.py

app.py

@@ -16,6 +16,15 @@ import requests
 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
+
 # 번역 모델 로드
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 
@@ -479,6 +488,31 @@ css = '''
 #component-11{align-self: stretch;}
 footer {visibility: hidden;}
 '''
+# 업스케일 함수 추가
+@spaces.GPU
+def upscale(input_image, progress=gr.Progress(track_tqdm=True)):
+    # 입력 이미지 처리
+    input_image, w_original, h_original, was_resized = process_input(input_image, 4)
+    
+    # 4096x4096 크기로 조정
+    control_image = input_image.resize((4096, 4096))
+
+    generator = torch.Generator(device=device).manual_seed(random.randint(0, MAX_SEED))
+
+    gr.Info("Upscaling image to 4096x4096...")
+    image = pipe(
+        prompt="",
+        control_image=control_image,
+        controlnet_conditioning_scale=0.6,
+        num_inference_steps=28,
+        guidance_scale=3.5,
+        height=4096,
+        width=4096,
+        generator=generator,
+    ).images[0]
+
+    return image
+    
 
 with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as app:
 
@@ -489,6 +523,7 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as a
             prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Type a prompt after selecting a LoRA")
         with gr.Column(scale=1):
             generate_button = gr.Button("Generate", variant="primary", elem_classes=["button_total"])
+            upscale_button = gr.Button("업스케일(4096X4096픽셀)", variant="secondary")
     with gr.Row(elem_id="loaded_loras"):
         with gr.Column(scale=1, min_width=25):
             randomize_button = gr.Button("🎲", variant="secondary", scale=1, elem_id="random_btn")
@@ -590,6 +625,12 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as a
         inputs=[result, history_gallery],
         outputs=history_gallery,
     )
+    
+    upscale_button.click(
+        upscale,
+        inputs=[result],
+        outputs=[result]
+    )
 
 app.queue()
 app.launch()