Update app.py

app.py

@@ -1,120 +1,210 @@
+import spaces
 import gradio as gr
 import torch
-import numpy as np
-from diffusers import DiffusionPipeline
-from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
-from functools import lru_cache
 from PIL import Image
-from transformers import CLIPImageProcessor
-
-@lru_cache(maxsize=1)
-def load_pipeline():
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    # Use FP16 when CUDA is available, along with a revision flag if supported.
-    torch_dtype = torch.float16 if device.type == "cuda" else torch.float32
-    revision = "fp16" if device.type == "cuda" else None
-
-    base_model = "black-forest-labs/FLUX.1-dev"
-    pipe = DiffusionPipeline.from_pretrained(
-        base_model,
-        torch_dtype=torch_dtype,
-        low_cpu_mem_usage=True,
-        revision=revision,
-    )
-    
-    # Load LoRA weights
-    lora_repo = "strangerzonehf/Flux-Super-Realism-LoRA"
-    pipe.load_lora_weights(lora_repo)
-    
-    # Load safety checker and image processor.
-    # If memory remains an issue, you can disable the safety checker below.
-    safety_checker = StableDiffusionSafetyChecker.from_pretrained(
-        "CompVis/stable-diffusion-safety-checker"
-    )
-    image_processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
-    
-    if device.type == "cuda":
-        # Use attention slicing for further memory savings.
-        pipe.enable_attention_slicing()
-        # Offload layers to CPU when not in use.
-        pipe.enable_sequential_cpu_offload()
-    
-    return pipe, safety_checker, image_processor
-
-pipe, safety_checker, image_processor = load_pipeline()
-
-def generate_image(
-    prompt,
-    seed=42,
-    width=512,   # Keep resolution low by default
-    height=512,
-    guidance_scale=6,
-    steps=28,
-    progress=gr.Progress()
+from diffusers import DiffusionPipeline
+import random
+import uuid
+from typing import Tuple
+import numpy as np
+
+def save_image(img):
+    unique_name = str(uuid.uuid4()) + ".png"
+    img.save(unique_name)
+    return unique_name
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+MAX_SEED = np.iinfo(np.int32).max
+
+if not torch.cuda.is_available():
+    DESCRIPTIONz += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"
+
+base_model = "black-forest-labs/FLUX.1-dev"
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
+
+lora_repo = "strangerzonehf/Flux-Super-Realism-LoRA"
+trigger_word = "Super Realism"  # Leave trigger_word blank if not used.
+
+pipe.load_lora_weights(lora_repo)
+pipe.to("cuda")
+
+style_list = [
+    {
+        "name": "3840 x 2160",
+        "prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+    },
+    {
+        "name": "2560 x 1440",
+        "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+    },
+    {
+        "name": "HD+",
+        "prompt": "hyper-realistic 2K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+    },
+    {
+        "name": "Style Zero",
+        "prompt": "{prompt}",
+    },
+]
+
+styles = {k["name"]: k["prompt"] for k in style_list}
+
+DEFAULT_STYLE_NAME = "3840 x 2160"
+STYLE_NAMES = list(styles.keys())
+
+def apply_style(style_name: str, positive: str) -> str:
+    return styles.get(style_name, styles[DEFAULT_STYLE_NAME]).replace("{prompt}", positive)
+
+@spaces.GPU(duration=60, enable_queue=True)
+def generate(
+    prompt: str,
+    seed: int = 0,
+    width: int = 1024,
+    height: int = 1024,
+    guidance_scale: float = 3,
+    randomize_seed: bool = False,
+    style_name: str = DEFAULT_STYLE_NAME,
+    progress=gr.Progress(track_tqdm=True),
 ):
-    try:
-        progress(0, desc="Initializing...")
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        generator = torch.Generator(device=device).manual_seed(seed)
-        
-        # Auto-add trigger words if not present
-        if "super realism" not in prompt.lower():
-            prompt = f"Super Realism, {prompt}"
-        
-        with torch.inference_mode():
-            result = pipe(
-                prompt=prompt,
-                width=width,
-                height=height,
-                guidance_scale=guidance_scale,
-                num_inference_steps=steps,
-                generator=generator,
-            )
-            image = result.images[0]
-        
-        progress(1, desc="Safety checking...")
-        # Process image for safety checking
-        safety_input = image_processor(image, return_tensors="pt")
-        np_image = np.array(image)
-        _, nsfw_detected = safety_checker(
-            images=[np_image],
-            clip_input=safety_input.pixel_values
-        )
-        
-        if nsfw_detected[0]:
-            return Image.new("RGB", (width, height)), "NSFW content detected"
-        
-        # Clear CUDA cache
-        if device.type == "cuda":
-            torch.cuda.empty_cache()
-        return image, "Generation successful"
-    
-    except Exception as e:
-        return Image.new("RGB", (width, height)), f"Error: {str(e)}"
+    seed = int(randomize_seed_fn(seed, randomize_seed))
 
-with gr.Blocks() as app:
-    gr.Markdown("# Flux Super Realism Generator")
+    positive_prompt = apply_style(style_name, prompt)
     
+    if trigger_word:
+        positive_prompt = f"{trigger_word} {positive_prompt}"
+
+    images = pipe(
+        prompt=positive_prompt,
+        width=width,
+        height=height,
+        guidance_scale=guidance_scale,
+        num_inference_steps=28,
+        num_images_per_prompt=1,
+        output_type="pil",
+    ).images
+    image_paths = [save_image(img) for img in images]
+    print(image_paths)
+    return image_paths, seed
+
+examples = [
+    "Woman in a red jacket, snowy, in the style of hyper-realistic portraiture, caninecore, mountainous vistas, timeless beauty, palewave, iconic, distinctive noses --ar 72:101 --stylize 750 --v 6",
+    "Super Realism, Headshot of handsome young man, wearing dark gray sweater with buttons and big shawl collar, brown hair and short beard, serious look on his face, black background, soft studio lighting, portrait photography --ar 85:128 --v 6.0 --style",
+    "Super Realism, High-resolution photograph, woman, UHD, photorealistic, shot on a Sony A7III --chaos 20 --ar 1:2 --style raw --stylize 250",
+    "Super-realism, Purple Dreamy, a medium-angle shot of a young woman with long brown hair, wearing a pair of eye-level glasses, stands in front of a backdrop of purple and white lights. The womans eyes are closed, her lips are slightly parted, as if she is looking up at the sky. Her hair is cascading over her shoulders, framing her face. She is wearing a sleeveless top, adorned with tiny white dots, and a gold chain necklace around her neck. Her left earrings are dangling from her ears, adding a pop of color to the scene."
+]
+
+css = '''
+.gradio-container{max-width: 888px !important}
+h1{text-align:center}
+footer {
+    visibility: hidden
+}
+.submit-btn {
+    background-color: #e34949 !important;
+    color: white !important;
+}
+.submit-btn:hover {
+    background-color: #ff3b3b !important;
+}
+'''
+
+with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     with gr.Row():
-        with gr.Column():
-            prompt_input = gr.Textbox(label="Prompt", value="A portrait of a person")
-            seed_input = gr.Slider(0, 1000, value=42, label="Seed")
-            # Limit the resolution sliders to help avoid memory overuse.
-            width_input = gr.Slider(256, 1024, value=512, step=64, label="Width")
-            height_input = gr.Slider(256, 1024, value=512, step=64, label="Height")
-            guidance_input = gr.Slider(1, 20, value=6, label="Guidance Scale")
-            steps_input = gr.Slider(10, 100, value=28, label="Steps")
-            submit = gr.Button("Generate")
+        with gr.Column(scale=1):
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=1,
+                placeholder="Enter your prompt",
+                container=False,
+            )
+            run_button = gr.Button("Generate as ( 768 x 1024 )🤗", scale=0, elem_classes="submit-btn")
+            
+            with gr.Accordion("Advanced options", open=True, visible=True):
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
+                    step=1,
+                    value=0,
+                    visible=True
+                )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                
+                with gr.Row(visible=True):
+                    width = gr.Slider(
+                        label="Width",
+                        minimum=512,
+                        maximum=2048,
+                        step=64,
+                        value=768,
+                    )
+                    height = gr.Slider(
+                        label="Height",
+                        minimum=512,
+                        maximum=2048,
+                        step=64,
+                        value=1024,
+                    )
+                
+                with gr.Row():
+                    guidance_scale = gr.Slider(
+                        label="Guidance Scale",
+                        minimum=0.1,
+                        maximum=20.0,
+                        step=0.1,
+                        value=3.0,
+                    )
+                    num_inference_steps = gr.Slider(
+                        label="Number of inference steps",
+                        minimum=1,
+                        maximum=40,
+                        step=1,
+                        value=28,
+                    )
+
+                style_selection = gr.Radio(
+                    show_label=True,
+                    container=True,
+                    interactive=True,
+                    choices=STYLE_NAMES,
+                    value=DEFAULT_STYLE_NAME,
+                    label="Quality Style",
+                )
         
-        with gr.Column():
-            output_image = gr.Image(label="Result", type="pil")
-            status = gr.Textbox(label="Status")
-    
-    submit.click(
-        generate_image,
-        inputs=[prompt_input, seed_input, width_input, height_input, guidance_input, steps_input],
-        outputs=[output_image, status]
+        with gr.Column(scale=2):
+            result = gr.Gallery(label="Result", columns=1, show_label=False)
+            
+            gr.Examples(
+                examples=examples,
+                inputs=prompt,
+                outputs=[result, seed],
+                fn=generate,
+                cache_examples=False,
+            )
+
+    gr.on(
+        triggers=[
+            prompt.submit,
+            run_button.click,
+        ],
+        fn=generate,
+        inputs=[
+            prompt,
+            seed,
+            width,
+            height,
+            guidance_scale,
+            randomize_seed,
+            style_selection,
+        ],
+        outputs=[result, seed],
+        api_name="run",
     )
-    
-    # Queue settings to limit concurrent requests
-    app.queue(max_size=3).launch()
+
+if __name__ == "__main__":
+    demo.queue(max_size=40).launch()