Spaces:
Runtime error
Runtime error
| import os | |
| import random | |
| import numpy as np | |
| from huggingface_hub import AsyncInferenceClient | |
| from translatepy import Translator | |
| from gradio_client import Client, handle_file | |
| from PIL import Image | |
| # Constants | |
| MAX_SEED = np.iinfo(np.int32).max | |
| HF_TOKEN = os.getenv('HF_TOKEN') # Set the environment variable for HF_TOKEN | |
| HF_TOKEN_UPSCALER = os.getenv('HF_TOKEN') # Set the environment variable for HF_TOKEN_UPSCALER | |
| class Lorify: | |
| def __init__(self, hf_token=None, hf_token_upscaler=None): | |
| # Optionally load tokens from environment if not passed | |
| self.hf_token = hf_token or HF_TOKEN | |
| self.hf_token_upscaler = hf_token_upscaler or HF_TOKEN_UPSCALER | |
| # Initialize clients | |
| self.qwen_client = Client("K00B404/HugChatWrap", hf_token=self.hf_token) | |
| self.client = AsyncInferenceClient() | |
| # List of available LoRAs (replace with your LoRA repo names or paths) | |
| self.loaded_loras = [] | |
| self.loras = [ | |
| "Shakker-Labs/FLUX.1-dev-LoRA-add-details", | |
| "XLabs-AI/flux-RealismLora", | |
| "enhanceaiteam/Flux-uncensored" | |
| ] | |
| self.loaded_loras.extend(self.loras) | |
| # Enable or disable LoRA | |
| def enable_lora(self, lora_add, basemodel): | |
| return basemodel if not lora_add else lora_add | |
| # Generate image function | |
| async def generate_image(self, prompt, model, lora_word, width, height, scales, steps, seed): | |
| try: | |
| if seed == -1: | |
| seed = random.randint(0, MAX_SEED) | |
| seed = int(seed) | |
| # Translate prompt | |
| text = str(Translator().translate(prompt, 'English')) + "," + lora_word | |
| # Generate image | |
| image = await self.client.text_to_image( | |
| prompt=text, | |
| height=height, | |
| width=width, | |
| guidance_scale=scales, | |
| num_inference_steps=steps, | |
| model=model | |
| ) | |
| return image, seed | |
| except Exception as e: | |
| print(f"Error generating image: {e}") | |
| return None, None | |
| # Upscale image function | |
| def upscale_image(self, prompt, img_path, upscale_factor): | |
| try: | |
| # Initialize the upscale client | |
| upscale_client = Client("finegrain/finegrain-image-enhancer", hf_token=self.hf_token_upscaler) | |
| result = upscale_client.predict( | |
| input_image=handle_file(img_path), | |
| prompt=prompt, | |
| negative_prompt="worst quality, low quality, normal quality", | |
| upscale_factor=upscale_factor, | |
| controlnet_scale=0.6, | |
| controlnet_decay=1, | |
| condition_scale=6, | |
| denoise_strength=0.35, | |
| num_inference_steps=18, | |
| solver="DDIM", | |
| api_name="/process" | |
| ) | |
| return result[1] # Return upscale image path | |
| except Exception as e: | |
| print(f"Error scaling image: {e}") | |
| return None | |
| # Main method to generate and optionally upscale image | |
| async def gen_image(self, prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora): | |
| model = self.enable_lora(lora_model, basemodel) if process_lora else basemodel | |
| image, seed = await self.generate_image(prompt, model, "", width, height, scales, steps, seed) | |
| if image is None: | |
| print("Image generation failed.") | |
| return [] | |
| image_path = "temp_image.jpg" | |
| image.save(image_path, format="JPEG") | |
| upscale_image_path = None | |
| if process_upscale: | |
| upscale_image_path = self.upscale_image(prompt, image_path, upscale_factor) | |
| if upscale_image_path and os.path.exists(upscale_image_path): | |
| return [image_path, upscale_image_path] | |
| return [image_path] |