Update app.py

app.py

@@ -1,27 +1,49 @@
 import gradio as gr
-import requests
-import io
+import torch
 import random
 import os
 import time
 from PIL import Image
 from deep_translator import GoogleTranslator
+from diffusers import DiffusionPipeline
+from huggingface_hub import hf_hub_download
 
-# Project by Nymbo
+# Project by Nymbo with LoRA integration
 
-API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
-API_TOKEN = os.getenv("HF_READ_TOKEN")
-headers = {"Authorization": f"Bearer {API_TOKEN}"}
-timeout = 100
+# Model and LoRA configuration
+BASE_MODEL = "black-forest-labs/FLUX.1-dev"
+LORA_REPO = "burhansyam/uncen"
+LORA_WEIGHTS_NAME = "uncen.safetensors"  # Adjust if different
+torch_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+
+# Initialize the pipeline with LoRA
+def init_pipeline():
+    pipe = DiffusionPipeline.from_pretrained(
+        BASE_MODEL,
+        torch_dtype=torch_dtype
+    )
+    
+    # Load LoRA weights
+    pipe.load_lora_weights(
+        hf_hub_download(repo_id=LORA_REPO, filename=LORA_WEIGHTS_NAME),
+        adapter_name="uncen"
+    )
+    
+    # Enable model offloading if needed
+    if torch.cuda.is_available():
+        pipe.to("cuda")
+        pipe.enable_xformers_memory_efficient_attention()
+    
+    return pipe
+
+pipe = init_pipeline()
 
 def convert_to_png(image):
     """Convert any image format to true PNG format"""
     png_buffer = io.BytesIO()
     if image.mode == 'RGBA':
-        # If image has alpha channel, save as PNG with transparency
         image.save(png_buffer, format='PNG', optimize=True)
     else:
-        # Convert to RGB first if not in RGB/RGBA mode
         if image.mode != 'RGB':
             image = image.convert('RGB')
         image.save(png_buffer, format='PNG', optimize=True)
@@ -34,8 +56,6 @@ def query(prompt, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Ka
         return None
 
     key = random.randint(0, 999)
-    API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN")])
-    headers = {"Authorization": f"Bearer {API_TOKEN}"}
     
     # Translate prompt
     try:
@@ -47,39 +67,47 @@ def query(prompt, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Ka
     
     print(f'\033[1mGeneration {key}:\033[0m {prompt}')
     
-    payload = {
-        "inputs": prompt,
-        "is_negative": is_negative,
-        "steps": steps,
-        "cfg_scale": cfg_scale,
-        "seed": seed if seed != -1 else random.randint(1, 1000000000),
-        "strength": strength,
-        "parameters": {"width": width, "height": height}
+    # Set random seed if not specified
+    generator = None
+    if seed != -1:
+        generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
+    else:
+        seed = random.randint(1, 1000000000)
+        generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
+    
+    # Map sampler names to Diffusers scheduler names
+    sampler_map = {
+        "DPM++ 2M Karras": "dpmsolver++",
+        "DPM++ SDE Karras": "dpmsolver++",
+        "Euler": "euler",
+        "Euler a": "euler_a",
+        "Heun": "heun",
+        "DDIM": "ddim"
     }
-
+    
     try:
-        response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
-        response.raise_for_status()
-        
-        # Convert directly to PNG without intermediate format
-        img = Image.open(io.BytesIO(response.content))
-        png_img = convert_to_png(img)
+        # Generate image with LoRA
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=is_negative if is_negative else None,
+            num_inference_steps=steps,
+            guidance_scale=cfg_scale,
+            generator=generator,
+            strength=strength,
+            width=width,
+            height=height,
+            cross_attention_kwargs={"scale": 0.8},  # LoRA strength adjustment
+        ).images[0]
         
+        png_img = convert_to_png(image)
         print(f'\033[1mGeneration {key} completed as PNG!\033[0m')
         return png_img
         
-    except requests.exceptions.RequestException as e:
-        print(f"API Error: {e}")
-        if hasattr(e, 'response') and e.response:
-            if e.response.status_code == 503:
-                raise gr.Error("503: Model is loading, please try again later")
-            raise gr.Error(f"{e.response.status_code}: {e.response.text}")
-        raise gr.Error("Network error occurred")
     except Exception as e:
-        print(f"Image processing error: {e}")
-        raise gr.Error(f"Image processing failed: {str(e)}")
+        print(f"Generation error: {e}")
+        raise gr.Error(f"Image generation failed: {str(e)}")
 
-# Light theme CSS
+# Light theme CSS (same as before)
 css = """
 #app-container {
     max-width: 800px;
@@ -111,7 +139,7 @@ h1 {
 """
 
 with gr.Blocks(theme=gr.themes.Default(primary_hue="green"), css=css) as app:
-    gr.HTML("<center><h1>FLUX.1-Dev (PNG Output)</h1></center>")
+    gr.HTML("<center><h1>FLUX.1-Dev with LoRA (PNG Output)</h1></center>")
     
     with gr.Column(elem_id="app-container"):
         with gr.Row():
@@ -151,7 +179,7 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue="green"), css=css) as app:
             output_image = gr.Image(
                 type="pil",
                 label="Generated PNG Image",
-                format="png",  # Explicitly set output format
+                format="png",
                 elem_id="gallery"
             )