Update app.py

app.py

@@ -22,6 +22,10 @@ from funcs import (
 )
 from transformers import pipeline
 
+import numpy as np
+from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+
 def download_model():
     REPO_ID = 'Doubiiu/DynamiCrafter_1024'
     filename_list = ['model.ckpt']
@@ -47,6 +51,30 @@ model = model.cuda()
 # 번역 모델 초기화
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 
+# FLUX 모델 설정
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to(device)
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 2048
+
+@spaces.GPU(duration=190)
+def infer_t2i(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed)
+    image = pipe(
+        prompt = prompt, 
+        width = width,
+        height = height,
+        num_inference_steps = num_inference_steps, 
+        generator = generator,
+        guidance_scale=guidance_scale
+    ).images[0] 
+    return image, seed
+
 @spaces.GPU(duration=300)
 def infer(image, prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, video_length=2):
     # 한글 입력 감지 및 번역
@@ -182,4 +210,28 @@ with gr.Blocks(analytics_enabled=False, css=css) as dynamicrafter_iface:
                         fn = infer
         )
 
+
+    with gr.Tab(label='T2I'):
+        with gr.Column():
+            with gr.Row():
+                t2i_input_text = gr.Text(label='Prompt')
+            with gr.Row():
+                t2i_seed = gr.Slider(label='Seed', minimum=0, maximum=MAX_SEED, step=1, value=42)
+                t2i_randomize_seed = gr.Checkbox(label='Randomize seed', value=False)
+            with gr.Row():
+                t2i_width = gr.Slider(label='Width', minimum=256, maximum=MAX_IMAGE_SIZE, step=64, value=1024)
+                t2i_height = gr.Slider(label='Height', minimum=256, maximum=MAX_IMAGE_SIZE, step=64, value=1024)
+            with gr.Row():
+                t2i_guidance_scale = gr.Slider(label='Guidance Scale', minimum=1.0, maximum=20.0, step=0.1, value=5.0)
+                t2i_num_inference_steps = gr.Slider(label='Inference Steps', minimum=1, maximum=100, step=1, value=28)
+            t2i_generate_btn = gr.Button("Generate")
+            t2i_output_image = gr.Image(label="Generated Image", elem_id="t2i_output_img")
+            t2i_output_seed = gr.Number(label="Used Seed", elem_id="t2i_output_seed")
+
+        t2i_generate_btn.click(
+            fn=infer_t2i,
+            inputs=[t2i_input_text, t2i_seed, t2i_randomize_seed, t2i_width, t2i_height, t2i_guidance_scale, t2i_num_inference_steps],
+            outputs=[t2i_output_image, t2i_output_seed]
+        )
+
 dynamicrafter_iface.queue(max_size=12).launch(show_api=True)