Update app.py

app.py

@@ -9,17 +9,19 @@ from custom_pipeline import FluxWithCFGPipeline
 
 # --- Torch Optimizations ---
 torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.benchmark = True # Enable cuDNN benchmark for potentially faster convolutions
+torch.backends.cudnn.benchmark = True
 
 # --- Constants ---
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 2048 # Keep a reasonable limit to prevent OOMs
+MAX_IMAGE_SIZE = 1024
 DEFAULT_WIDTH = 1024
-DEFAULT_HEIGHT = 1024
-DEFAULT_INFERENCE_STEPS = 1 # FLUX Schnell is designed for few steps
-MIN_INFERENCE_STEPS = 1
-MAX_INFERENCE_STEPS = 8 # Allow slightly more steps for potential quality boost
-ENHANCE_STEPS = 2 # Fixed steps for the enhance button
+DEFAULT_HEIGHT = 576
+ASPECT_RATIOS = {
+    "16:9": (1024, 576),
+    "1:1": (1024, 1024),
+    "9:16": (576, 1024)
+}
+INFERENCE_STEPS = 8
 
 # --- Device and Model Setup ---
 dtype = torch.float16
@@ -27,174 +29,95 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 pipe = FluxWithCFGPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype)
 pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
-
 pipe.to(device)
 
 # --- Inference Function ---
 @spaces.GPU
-def generate_image(prompt: str, seed: int = 42, width: int = DEFAULT_WIDTH, height: int = DEFAULT_HEIGHT, randomize_seed: bool = False, num_inference_steps: int = DEFAULT_INFERENCE_STEPS, is_enhance: bool = False):
-    """Generates an image using the FLUX pipeline with error handling."""
-
+def generate_image(prompt: str, seed: int = 42, aspect_ratio: str = "16:9", randomize_seed: bool = False):
     if pipe is None:
-        raise gr.Error("Diffusion pipeline failed to load. Cannot generate images.")
-        
+        raise gr.Error("Pipelinei nuk u ngarkua.")
+
     if not prompt or prompt.strip() == "":
-         gr.Warning("Prompt is empty. Please enter a description.")
-         return None, seed, "Error: Empty prompt"
+        return None, seed, "Gabim: Plotësoni përshkrimin."
 
-    start_time = time.time()
-    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    
-    # Clamp dimensions to avoid excessive memory usage
-    width = min(width, MAX_IMAGE_SIZE)
-    height = min(height, MAX_IMAGE_SIZE)
-    
-    # Use fixed steps for enhance button, otherwise use slider value
-    steps_to_use = ENHANCE_STEPS if is_enhance else num_inference_steps
-    # Clamp steps
-    steps_to_use = max(MIN_INFERENCE_STEPS, min(steps_to_use, MAX_INFERENCE_STEPS))
+
+    width, height = ASPECT_RATIOS.get(aspect_ratio, (DEFAULT_WIDTH, DEFAULT_HEIGHT))
+
+    # Allow prompts in Albanian or English and still enhance them
+    enhanced_prompt = prompt.strip() + ", ultra realistic, sharp, 8k resolution, cinematic lighting"
 
     try:
-        # Ensure generator is on the correct device
         generator = torch.Generator(device=device).manual_seed(int(float(seed)))
-
-        # Use inference_mode for efficiency
+        start_time = time.time()
         with torch.inference_mode():
-            # Generate the image (assuming pipe returns list/tuple with image first)
-            # Modify pipe call based on its actual signature if needed
-            result_img = pipe(
-                prompt=prompt,
+            image = pipe(
+                prompt=enhanced_prompt,
                 width=width,
                 height=height,
-                num_inference_steps=steps_to_use,
+                num_inference_steps=INFERENCE_STEPS,
                 generator=generator,
-                output_type="pil", # Ensure PIL output for Gradio Image component
-                return_dict=False # Assuming the custom pipeline supports this for direct output
-            )[0][0] # Assuming the output structure is [[img]]
-
+                output_type="pil",
+                return_dict=False
+            )[0][0]
         latency = time.time() - start_time
-        latency_str = f"Latency: {latency:.2f} seconds (Steps: {steps_to_use})"
-        return result_img, seed, latency_str
-
-    except torch.cuda.OutOfMemoryError as e:
-        # Clear cache and suggest reducing size/steps
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        raise gr.Error("GPU ran out of memory. Try reducing the image width/height or the number of inference steps.")
-        
+        return image, seed, f"Koha e përpunimit: {latency:.2f} sekonda"
     except Exception as e:
-        # Clear cache just in case
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
-        raise gr.Error(f"An error occurred during generation: {e}")
+        raise gr.Error(f"Gabim gjatë gjenerimit: {e}")
 
-
-# --- Example Prompts ---
 examples = [
-    "a tiny astronaut hatching from an egg on the moon",
-    "a cute white cat holding a sign that says hello world",
-    "an anime illustration of Steve Jobs",
-    "Create image of Modern house in minecraft style",
-    "photo of a woman on the beach, shot from above. She is facing the sea, while wearing a white dress. She has long blonde hair",
-    "Selfie photo of a wizard with long beard and purple robes, he is apparently in the middle of Tokyo. Probably taken from a phone.",
-    "Photo of a young woman with long, wavy brown hair tied in a bun and glasses. She has a fair complexion and is wearing subtle makeup, emphasizing her eyes and lips. She is dressed in a black top. The background appears to be an urban setting with a building facade, and the sunlight casts a warm glow on her face.",
-    "High-resolution photorealistic render of a sleek, futuristic motorcycle parked on a neon-lit street at night, rain reflecting the lights.",
-    "Watercolor painting of a cozy bookstore interior with overflowing shelves and a cat sleeping in a sunbeam.",
+    "Qytet futuristik natën me drita neon",
+    "Një mace e bardhë që mban një tabelë përshëndetëse",
+    "Një astronaut që del nga një vezë në Hënë",
+    "Pamje nga një shtëpi moderne në stilin Minecraft"
 ]
 
-# --- Gradio UI ---
-with gr.Blocks() as demo:
-    with gr.Column(elem_id="app-container"):
-        gr.Markdown("# 🎨 Realtime FLUX Image Generator")
-        gr.Markdown("Generate stunning images in real-time with Modified Flux.Schnell pipeline.")
-        gr.Markdown("<span style='color: red;'>Note: Sometimes it stucks or stops generating images (I don't know why). In that situation just refresh the site.</span>")
-
-        with gr.Row():
-            with gr.Column(scale=2.5):
-                result = gr.Image(label="Generated Image", show_label=False, interactive=False)
-            with gr.Column(scale=1):
-                prompt = gr.Text(
-                    label="Prompt",
-                    placeholder="Describe the image you want to generate...",
-                    lines=3,
-                    show_label=False,
-                    container=False,
-                )
-                generateBtn = gr.Button("🖼️ Generate Image")
-                enhanceBtn = gr.Button("🚀 Enhance Image")
-
-                with gr.Column("Advanced Options"):
-                    with gr.Row():
-                        realtime = gr.Checkbox(label="Realtime Toggler", info="If TRUE then uses more GPU but create image in realtime.", value=False)
-                        latency = gr.Text(label="Latency")
-                    with gr.Row():
-                        seed = gr.Number(label="Seed", value=42)
-                        randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-                    with gr.Row():
-                        width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_WIDTH)
-                        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_HEIGHT)
-                        num_inference_steps = gr.Slider(label="Inference Steps", minimum=MIN_INFERENCE_STEPS, maximum=MAX_INFERENCE_STEPS, step=1, value=DEFAULT_INFERENCE_STEPS)
-
-        with gr.Row():
-            gr.Markdown("### 🌟 Inspiration Gallery")
-        with gr.Row():
-            gr.Examples(
-                examples=examples,
-                fn=generate_image,
-                inputs=[prompt],
-                outputs=[result, seed, latency],
-                cache_examples=True,
-                cache_mode="eager"
-            )
-
-    enhanceBtn.click(
-        fn=generate_image,
-        inputs=[prompt, seed, width, height],
-        outputs=[result, seed, latency],
-        show_progress="full"
-    )
-
-    generateBtn.click(
+# --- App Layout ---
+with gr.Blocks(css="""
+body::before {
+    content: "";
+    display: block;
+    height: 640px;
+    background-color: #0f1117;
+}
+button[aria-label="Download"] {
+    transform: scale(1.5);
+    transform-origin: top right;
+    margin: 0 !important;
+    padding: 6px !important;
+}
+""") as app:
+    gr.Markdown("# 🖼️ Gjenerues Imazhesh FLUX")
+    gr.Markdown("Përdor modelin FLUX për të krijuar imazhe fantastike nga përshkrime në **gjuhën shqipe ose angleze**.")
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            output_image = gr.Image(label="Imazhi i Gjeneruar", interactive=False, show_download_button=True)
+        with gr.Column(scale=1):
+            prompt = gr.Text(label="Përshkrimi", placeholder="Shkruani se çfarë doni të krijoni...", lines=3)
+            generate_btn = gr.Button("🎨 Gjenero")
+            aspect_ratio = gr.Radio(label="Raporti i Imazhit", choices=list(ASPECT_RATIOS.keys()), value="16:9")
+            randomize_seed = gr.Checkbox(label="Përdor numër të rastësishëm", value=True)
+            latency = gr.Text(label="Koha", interactive=False)
+
+    gr.Markdown("### 📌 Shembuj Frymëzues")
+    gr.Examples(
+        examples=examples,
         fn=generate_image,
-        inputs=[prompt, seed, width, height, randomize_seed, num_inference_steps],
-        outputs=[result, seed, latency],
-        show_progress="full",
-        api_name="RealtimeFlux",
-    )
-
-    def update_ui(realtime_enabled):
-        return {
-            prompt: gr.update(interactive=True),
-            generateBtn: gr.update(visible=not realtime_enabled)
-        }
-
-    def realtime_generation(*args):
-        if args[0]:  # If realtime is enabled
-            return next(generate_image(*args[1:]))
-
-    realtime.change(
-        fn=update_ui,
-        inputs=[realtime],
-        outputs=[prompt, generateBtn]
+        inputs=[prompt],
+        outputs=[output_image, gr.Number(visible=False), latency],
+        cache_examples=True,
+        cache_mode="eager"
     )
 
-    prompt.submit(
+    generate_btn.click(
         fn=generate_image,
-        inputs=[prompt, seed, width, height, randomize_seed, num_inference_steps],
-        outputs=[result, seed, latency],
+        inputs=[prompt, gr.Number(value=42, visible=False), aspect_ratio, randomize_seed],
+        outputs=[output_image, gr.Number(visible=False), latency],
         show_progress="full"
     )
 
-    for component in [prompt, width, height, num_inference_steps]:
-        component.input(
-            fn=realtime_generation,
-            inputs=[realtime, prompt, seed, width, height, randomize_seed, num_inference_steps],
-            outputs=[result, seed, latency],
-            show_progress="hidden",
-            trigger_mode="always_last"
-        )
-
-# Launch the app
-demo.launch()
+app.launch(share=True)