Upload 2 files

app (28).py

@@ -0,0 +1,295 @@
+import spaces
+import argparse
+import os
+import time
+from os import path
+import shutil
+from datetime import datetime
+from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download
+import gradio as gr
+import torch
+from diffusers import FluxPipeline
+from PIL import Image
+from transformers import pipeline
+
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+
+# Hugging Face 토큰 설정
+HF_TOKEN = os.getenv("HF_TOKEN")
+if HF_TOKEN is None:
+    raise ValueError("HF_TOKEN environment variable is not set")
+
+# Setup and initialization code
+cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
+PERSISTENT_DIR = os.environ.get("PERSISTENT_DIR", ".")
+gallery_path = path.join(PERSISTENT_DIR, "gallery")
+
+os.environ["TRANSFORMERS_CACHE"] = cache_path
+os.environ["HF_HUB_CACHE"] = cache_path
+os.environ["HF_HOME"] = cache_path
+
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# Create gallery directory if it doesn't exist
+if not path.exists(gallery_path):
+    os.makedirs(gallery_path, exist_ok=True)
+
+class timer:
+    def __init__(self, method_name="timed process"):
+        self.method = method_name
+    def __enter__(self):
+        self.start = time.time()
+        print(f"{self.method} starts")
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        end = time.time()
+        print(f"{self.method} took {str(round(end - self.start, 2))}s")
+
+# Model initialization
+if not path.exists(cache_path):
+    os.makedirs(cache_path, exist_ok=True)
+
+# 인증된 모델 로드
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    torch_dtype=torch.bfloat16,
+    use_auth_token=HF_TOKEN
+)
+
+# Hyper-SD LoRA 로드
+pipe.load_lora_weights(
+    hf_hub_download(
+        "ByteDance/Hyper-SD",
+        "Hyper-FLUX.1-dev-8steps-lora.safetensors",
+        use_auth_token=HF_TOKEN
+    )
+)
+pipe.fuse_lora(lora_scale=0.125)
+pipe.to(device="cuda", dtype=torch.bfloat16)
+
+def save_image(image):
+    """Save the generated image and return the path"""
+    try:
+        if not os.path.exists(gallery_path):
+            try:
+                os.makedirs(gallery_path, exist_ok=True)
+            except Exception as e:
+                print(f"Failed to create gallery directory: {str(e)}")
+                return None
+        
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        random_suffix = os.urandom(4).hex()
+        filename = f"generated_{timestamp}_{random_suffix}.png"
+        filepath = os.path.join(gallery_path, filename)
+        
+        try:
+            if isinstance(image, Image.Image):
+                image.save(filepath, "PNG", quality=100)
+            else:
+                image = Image.fromarray(image)
+                image.save(filepath, "PNG", quality=100)
+            
+            return filepath
+        except Exception as e:
+            print(f"Failed to save image: {str(e)}")
+            return None
+            
+    except Exception as e:
+        print(f"Error in save_image: {str(e)}")
+        return None
+
+# 예시 프롬프트 정의
+examples = [
+    ["A 3D Star Wars Darth Vader helmet, highly detailed metallic finish"],
+    ["A 3D Iron Man mask with glowing eyes and metallic red-gold finish"],
+    ["A detailed 3D Pokemon Pikachu figure with glossy surface"],
+    ["A 3D geometric abstract cube transforming into a sphere, metallic finish"],
+    ["A 3D steampunk mechanical heart with brass and copper details"],
+    ["A 3D crystal dragon with transparent iridescent scales"],
+    ["A 3D futuristic hovering drone with neon light accents"],
+    ["A 3D ancient Greek warrior helmet with ornate details"],
+    ["A 3D robotic butterfly with mechanical wings and metallic finish"],
+    ["A 3D floating magical crystal orb with internal energy swirls"]
+]
+
+@spaces.GPU
+def process_and_save_image(height=1024, width=1024, steps=8, scales=3.5, prompt="", seed=None):
+    global pipe
+    
+    if seed is None:
+        seed = torch.randint(0, 1000000, (1,)).item()
+    
+    # 한글 감지 및 번역
+    def contains_korean(text):
+        return any(ord('가') <= ord(c) <= ord('힣') for c in text)
+    
+    # 프롬프트 전처리
+    if contains_korean(prompt):
+        translated = translator(prompt)[0]['translation_text']
+        prompt = translated
+    
+    formatted_prompt = f"wbgmsst, 3D, {prompt} ,white background"
+    
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
+        try:
+            generated_image = pipe(
+                prompt=[formatted_prompt],
+                generator=torch.Generator().manual_seed(int(seed)),
+                num_inference_steps=int(steps),
+                guidance_scale=float(scales),
+                height=int(height),
+                width=int(width),
+                max_sequence_length=256
+            ).images[0]
+            
+            saved_path = save_image(generated_image)
+            if saved_path is None:
+                print("Warning: Failed to save generated image")
+            
+            return generated_image
+        except Exception as e:
+            print(f"Error in image generation: {str(e)}")
+            return None
+
+def get_random_seed():
+    return torch.randint(0, 1000000, (1,)).item()
+
+    
+def process_example(prompt):
+    return process_and_save_image(
+        height=1024,
+        width=1024,
+        steps=8,
+        scales=3.5,
+        prompt=prompt,
+        seed=get_random_seed()
+    )
+
+    
+# Gradio 인터페이스
+with gr.Blocks(
+    theme=gr.themes.Soft(),
+    css="""
+        .container {
+            background: linear-gradient(to bottom right, #1a1a1a, #4a4a4a);
+            border-radius: 20px;
+            padding: 20px;
+        }
+        .generate-btn {
+            background: linear-gradient(45deg, #2196F3, #00BCD4);
+            border: none;
+            color: white;
+            font-weight: bold;
+            border-radius: 10px;
+        }
+        .output-image {
+            border-radius: 15px;
+            box-shadow: 0 8px 16px rgba(0,0,0,0.2);
+        }
+        .fixed-width {
+            max-width: 1024px;
+            margin: auto;
+        }
+    """
+) as demo:
+    gr.HTML(
+        """
+        <div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">
+            <h1 style="font-size: 2.5rem; color: #2196F3;">3D Style Image Generator</h1>
+            <p style="font-size: 1.2rem; color: #666;">Create amazing 3D-style images with AI</p>
+        </div>
+        """
+    )
+    
+    with gr.Row(elem_classes="container"):
+        with gr.Column(scale=3):
+            prompt = gr.Textbox(
+                label="Image Description",
+                placeholder="Describe the 3D image you want to create...",
+                lines=3
+            )
+            
+            with gr.Accordion("Advanced Settings", open=False):
+                with gr.Row():
+                    height = gr.Slider(
+                        label="Height",
+                        minimum=256,
+                        maximum=1152,
+                        step=64,
+                        value=1024
+                    )
+                    width = gr.Slider(
+                        label="Width",
+                        minimum=256,
+                        maximum=1152,
+                        step=64,
+                        value=1024
+                    )
+                
+                with gr.Row():
+                    steps = gr.Slider(
+                        label="Inference Steps",
+                        minimum=6,
+                        maximum=25,
+                        step=1,
+                        value=8
+                    )
+                    scales = gr.Slider(
+                        label="Guidance Scale",
+                        minimum=0.0,
+                        maximum=5.0,
+                        step=0.1,
+                        value=3.5
+                    )
+                
+                seed = gr.Number(
+                    label="Seed (random by default, set for reproducibility)",
+                    value=get_random_seed(),
+                    precision=0
+                )
+                
+                randomize_seed = gr.Button("🎲 Randomize Seed", elem_classes=["generate-btn"])
+            
+            generate_btn = gr.Button(
+                "✨ Generate Image",
+                elem_classes=["generate-btn"]
+            )
+
+        with gr.Column(scale=4, elem_classes=["fixed-width"]):
+            output = gr.Image(
+                label="Generated Image",
+                elem_id="output-image",
+                elem_classes=["output-image", "fixed-width"],
+                value="3d.webp"
+            )
+
+    # Examples 섹션
+    gr.Examples(
+        examples=examples,
+        inputs=prompt,
+        outputs=output,
+        fn=process_example,  # 수정된 함수 사용
+        cache_examples=False,
+        examples_per_page=5
+    )
+
+    def update_seed():
+        return get_random_seed()
+
+    # 이벤트 핸들러
+    generate_btn.click(
+        process_and_save_image,
+        inputs=[height, width, steps, scales, prompt, seed],
+        outputs=output
+    ).then(
+        update_seed,
+        outputs=[seed]
+    )
+
+    randomize_seed.click(
+        update_seed,
+        outputs=[seed]
+    )
+
+if __name__ == "__main__":
+    demo.launch(allowed_paths=[PERSISTENT_DIR])

requirements (9).txt

@@ -0,0 +1,14 @@
+accelerate
+diffusers==0.30.0
+invisible_watermark
+torch
+transformers==4.43.3
+xformers
+sentencepiece
+peft
+
+safetensors
+gradio>=4.4.0
+Pillow>=9.0.0
+huggingface-hub>=0.19.0
+sacremoses