Update app.py

app.py

@@ -1,24 +1,61 @@
-import os
-import spaces
-import gradio as gr
-import numpy as np
-from PIL import Image
-import random
-from diffusers import DiffusionPipeline, AutoPipelineForText2Image
-from transformers import pipeline as transformers_pipeline
-import re
-from cohere import ClientV2
-import sys
 import torch
+from diffusers import AutoPipelineForText2Image
+from peft import PeftModel, PeftConfig
+import gradio as gr # Import Gradio
 
-import ast  #추가 삽입, requirements: albumentations 추가
-script_repr = os.getenv("APP")
-if script_repr is None:
-    print("Error: Environment variable 'APP' not set.")
-    sys.exit(1)
+# 기기 설정
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-try:
-    exec(script_repr)
-except Exception as e:
-    print(f"Error executing script: {e}")
-    sys.exit(1)
+# 기본 모델 로드
+print("기본 FLUX 모델 로드 중...")
+pipe = AutoPipelineForText2Image.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    torch_dtype=torch.float16  # bfloat16 대신 float16 사용
+)
+pipe.to(device)
+
+# Uncensored LoRA 로드
+print("Uncensored LoRA 로드 중...")
+pipe.load_lora_weights(
+    'Heartsync/Flux-NSFW-uncensored',
+    weight_name='lora.safetensors',
+    adapter_name="uncensored"
+)
+
+# 이미지 생성 함수 정의
+def generate_image(prompt, negative_prompt, guidance_scale, num_inference_steps, width, height, seed):
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    image = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        width=width,
+        height=height,
+        generator=generator,
+    ).images[0]
+    
+    # 이미지 저장 (선택 사항, 필요에 따라 활성화/비활성화)
+    # image.save("generated_image.png")
+    return image
+
+# Gradio 인터페이스 생성
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=[
+        gr.Textbox(label="Prompt", value="A woman in a sheer white dress standing on a beach at sunset, backlit so her silhouette is visible through the thin fabric, shot with Canon EOS R5, 85mm f/1.2 lens, golden hour natural lighting, professional composition, hyperrealistic detail, masterpiece quality, 8K resolution."),
+        gr.Textbox(label="Negative Prompt", value="text, watermark, signature, cartoon, anime, illustration, painting, drawing, low quality, blurry"),
+        gr.Slider(minimum=1.0, maximum=20.0, step=0.1, value=7.0, label="Guidance Scale"),
+        gr.Slider(minimum=10, maximum=100, step=1, value=28, label="Number of Inference Steps"),
+        gr.Slider(minimum=256, maximum=1024, step=64, value=1024, label="Width"),
+        gr.Slider(minimum=256, maximum=1024, step=64, value=1024, label="Height"),
+        gr.Slider(minimum=0, maximum=99999, step=1, value=42, label="Seed")
+    ],
+    outputs="image",
+    title="FLUX.1-dev with Uncensored LoRA",
+    description="Generate images using FLUX.1-dev with a loaded Uncensored LoRA model."
+)
+
+# 인터페이스 실행 (share=True)
+iface.launch(share=True)